---
title: Astract-1
canonical_url: "https://www.modelscope.cn/models/Isaac10086/Astract-1"
md_url: "https://www.modelscope.cn/models/Isaac10086/Astract-1.md"
repository: Isaac10086/Astract-1
chinese_name: "抽象-1"
last_updated: 2025-02-26
license: "Apache License 2.0"
pipeline_tag: text-generation
tasks:
  - text-generation
library_name:
  - pytorch
  - transformer
frameworks:
  - Pytorch
language:
  - zh
downloads: 8
stars: 0
---

# Astract-1

> Astract-1 - Isaac10086 在 ModelScope 开源的模型。110M左右参数，语料十分抽象

Isaac10086/Astract-1 是 ModelScope 魔搭社区上的text-generation模型，采用 Apache License 2.0 许可。

- **Repository**: Isaac10086/Astract-1
- **License**: Apache License 2.0
- **Tasks**: text-generation
- **Downloads**: 8
- **Stars**: 0
- **Last updated**: 2025-02-26

Source: https://www.modelscope.cn/models/Isaac10086/Astract-1

---

### 当前模型的贡献者未提供更加详细的模型介绍。模型文件和权重，可浏览“模型文件”页面获取。
#### 您可以通过如下git clone命令，或者ModelScope SDK来下载模型

SDK下载
```bash
#安装ModelScope
pip install modelscope
```
```python
#SDK模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('Isaac10086/Astract-1')
```
Git下载
```
#Git模型下载
git clone https://www.modelscope.cn/Isaac10086/Astract-1.git
```

<p style="color: lightgrey;">如果您是本模型的贡献者，我们邀请您根据<a href="https://modelscope.cn/docs/ModelScope%E6%A8%A1%E5%9E%8B%E6%8E%A5%E5%85%A5%E6%B5%81%E7%A8%8B%E6%A6%82%E8%A7%88" style="color: lightgrey; text-decoration: underline;">模型贡献文档</a>，及时完善模型卡片内容。</p>

```
import os
import platform
import argparse
import time
import math
import warnings
import pandas as pd
import torch
import torch.distributed as dist
from torch import optim, nn
from torch.nn.parallel import DistributedDataParallel
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader, DistributedSampler
from contextlib import nullcontext

from transformers import AutoTokenizer

from model.model import MiniMindLM
from model.LMConfig import LMConfig
from model.dataset import PretrainDataset

warnings.filterwarnings('ignore')

# 自定义日志函数（支持分布式训练）
def Logger(content):
    if not ddp or dist.get_rank() == 0:
        print(content)

# 动态学习率函数（含Warmup和分层学习率）
def get_lr(current_step, total_steps, base_lr, optimizer, warmup_steps=2000):
    # Warmup阶段
    if current_step < warmup_steps:
        lr = base_lr * current_step / warmup_steps
    else:
        # 余弦退火
        progress = (current_step - warmup_steps) / (total_steps - warmup_steps)
        lr = 0.5 * base_lr * (1 + math.cos(math.pi * progress))
    
    # 分层学习率（嵌入层学习率降低10倍）
    for param_group in optimizer.param_groups:
        if 'embed' in param_group['name']:
            param_group['lr'] = lr * 0.1
        else:
            param_group['lr'] = lr
    return lr

# 训练单个epoch
def train_epoch(epoch, wandb, best_val_loss, early_stop_counter):
    loss_fct = nn.CrossEntropyLoss(reduction='none')
    start_time = time.time()
    model.train()
    
    for step, (X, Y, loss_mask) in enumerate(train_loader):
        X, Y, loss_mask = X.to(args.device), Y.to(args.device), loss_mask.to(args.device)

        # 动态学习率更新
        current_step = epoch * iter_per_epoch + step
        lr = get_lr(current_step, args.epochs * iter_per_epoch, args.learning_rate, optimizer, args.warmup_iters)

        # 混合精度训练上下文
        with ctx:
            res = model(X)
            # 计算主损失
            loss = loss_fct(res.logits.view(-1, res.logits.size(-1)), Y.view(-1)).view(Y.size())
            loss = (loss * loss_mask).sum() / loss_mask.sum()
            # 添加MoE辅助损失（如果启用）
            if hasattr(res, 'aux_loss'):
                loss += res.aux_loss
            loss = loss / args.accumulation_steps

        # 反向传播与梯度累积
        scaler.scale(loss).backward()

        # 梯度裁剪与参数更新
        if (step + 1) % args.accumulation_steps == 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)

        # 日志记录
        if step % args.log_interval == 0:
            spend_time = time.time() - start_time
            Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iter_per_epoch}) '
                   f'loss:{loss.item()*args.accumulation_steps:.3f} '
                   f'lr:{optimizer.param_groups[0]["lr"]:.6f} '
                   f'epoch_Time:{spend_time/(step+1)*iter_per_epoch//60 - spend_time//60}min')
            
            if wandb and (not ddp or dist.get_rank() == 0):
                wandb.log({
                    "loss": loss.item() * args.accumulation_steps,
                    "lr": optimizer.param_groups[0]['lr']
                })

        # 模型保存与验证
        if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
            val_loss = validate()
            Logger(f"Validation Loss: {val_loss:.3f}")
            
            # 早停机制
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                early_stop_counter = 0
                save_model()
            else:
                early_stop_counter += 1
                if early_stop_counter >= args.early_stop_patience:
                    Logger("Early stopping triggered!")
                    return best_val_loss, True  # 停止训练
            
    return best_val_loss, False

# 验证函数
def validate():
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for X, Y, loss_mask in val_loader:
            X, Y, loss_mask = X.to(args.device), Y.to(args.device), loss_mask.to(args.device)
            res = model(X)
            loss = (nn.CrossEntropyLoss(reduction='none')(res.logits.view(-1, res.logits.size(-1)), 
                                                         Y.view(-1)).view(Y.size()) * loss_mask).sum() / loss_mask.sum()
            total_loss += loss.item()
    model.train()
    return total_loss / len(val_loader)

# 模型保存
def save_model():
    moe_path = '_moe' if lm_config.use_moe else ''
    ckp = f'{args.save_dir}/pretrain_{lm_config.dim}{moe_path}_best.pth'
    state_dict = model.module.state_dict() if ddp else model.state_dict()
    torch.save(state_dict, ckp)
    Logger(f"Model saved to {ckp}")

# 模型初始化
def init_model(lm_config):
    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
    model = MiniMindLM(lm_config).to(args.device)
    
    # 参数统计（验证参数量不变）
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    Logger(f'模型总参数量：{total_params / 1e6:.2f}M')
    
    return model, tokenizer

# 分布式初始化
def init_distributed_mode():
    if not ddp: return
    global ddp_local_rank, DEVICE

    dist.init_process_group(backend="nccl")
    ddp_rank = int(os.environ["RANK"])
    ddp_local_rank = int(os.environ["LOCAL_RANK"])
    ddp_world_size = int(os.environ["WORLD_SIZE"])
    DEVICE = f"cuda:{ddp_local_rank}"
    torch.cuda.set_device(DEVICE)

if __name__ == "__main__":
    # 参数配置（新增验证集和早停参数）
    parser = argparse.ArgumentParser(description="MiniMind Pretraining")
    parser.add_argument('--dim', default=1024, type=int)
    parser.add_argument('--n_layers', default=10, type=int)
    parser.add_argument('--val_data_path', type=str, required=True, help="验证集路径")
    parser.add_argument('--early_stop_patience', type=int, default=3, help="早停耐心值")
    # ...（其他参数保持不变，新增参数插入此处）...
    parser.add_argument("--out_dir", type=str, default="out")
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--learning_rate", type=float, default=5e-4)
    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--dtype", type=str, default="bfloat16")
    parser.add_argument("--use_wandb", action="store_true")
    parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain")
    parser.add_argument("--num_workers", type=int, default=1)
    parser.add_argument("--ddp", action="store_true")
    parser.add_argument("--accumulation_steps", type=int, default=8)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--warmup_iters", type=int, default=0)
    parser.add_argument("--log_interval", type=int, default=100)
    parser.add_argument("--save_interval", type=int, default=100)
    parser.add_argument('--local_rank', type=int, default=-1)
    parser.add_argument('--max_seq_len', default=512, type=int)
    parser.add_argument('--use_moe', default=False, type=bool)
    parser.add_argument("--data_path", type=str, default="./dataset/tfp.jsonl")
    args = parser.parse_args()

    # 模型配置
    lm_config = LMConfig(
        dim=args.dim,
        n_layers=args.n_layers,
        max_seq_len=args.max_seq_len,
        use_moe=args.use_moe,
        dropout=0.1  # 新增dropout配置
    )

    # 分布式初始化
    ddp = int(os.environ.get("RANK", -1)) != -1
    if ddp:
        init_distributed_mode()
        args.device = torch.device(DEVICE)

    # 混合精度与硬件优化
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision('high')  # 启用TF32加速

    # 初始化模型与数据
    model, tokenizer = init_model(lm_config)
    train_ds = PretrainDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
    val_ds = PretrainDataset(args.val_data_path, tokenizer, max_length=lm_config.max_seq_len)
    
    # 数据加载器
    train_sampler = DistributedSampler(train_ds) if ddp else None
    train_loader = DataLoader(
        train_ds,
        batch_size=args.batch_size,
        sampler=train_sampler,
        num_workers=args.num_workers,
        pin_memory=True
    )
    val_loader = DataLoader(val_ds, batch_size=args.batch_size)

    # 优化器（分层参数组）
    optimizer = optim.AdamW(
        [
            {"params": model.embed.parameters(), "name": "embed"},
            {"params": [p for n,p in model.named_parameters() if 'embed' not in n], "name": "transformer"}
        ],
        lr=args.learning_rate
    )

    # 分布式训练包装
    if ddp:
        model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
        model = DistributedDataParallel(model, device_ids=[ddp_local_rank])

    # 训练循环（含早停）
    scaler = torch.cuda.amp.GradScaler()
    iter_per_epoch = len(train_loader)
    best_val_loss = float('inf')
    early_stop_counter = 0
    early_stop = False

    for epoch in range(args.epochs):
        if early_stop:
            break
        best_val_loss, early_stop = train_epoch(epoch, None, best_val_loss, early_stop_counter)

    Logger("Training completed!")
```