这是使用 Llama 架构重现 roneneldan/TinyStories-1M 的第一个版本。

• 完整的训练过程包含在 train.ipynb 笔记本中。要重现此过程，只需将 TinyStoriesV2-GPT4-train.txt 和 TinyStoriesV2-GPT4-valid.txt 下载到与笔记本相同的文件夹中，然后运行各个单元格即可。脚本不会使用验证内容，因此您可以在其中放入任何内容。

• backup 目录中有一个 do_backup 脚本，我用它将权重从远程机器复制到本地。权重生成速度非常快，因此当脚本复制权重 N+1 时...

• 这是一个极其初级的概念验证（PoC）版本。训练会截断长于上下文大小的故事，并且不会使用任何滑动窗口来从故事的非起始位置进行训练。

• 在 40GB A100 上，训练大约需要 9 小时（每个 epoch 3 小时）。使用了约 30GB 的显存。

• 我使用的是 open_llama_3b 的分词器。不过，我在本地使用时遇到了问题（https://github.com/openlm-research/open_llama/issues/69）。在预装了库的云服务器上则没有问题。

• 演示脚本为 demo.py。

• 提供了验证脚本：valid.py。使用方法如下：python valid.py path/to/TinyStoriesV2-GPT4-valid.txt [optional-model-id-or-path]。训练后我认为没有必要将验证集分割成块。

• 此外，此版本使用了非常简单的缓存机制来对训练故事进行洗牌：它会保留最近加载的 N 个块的缓存，因此如果随机洗牌需要某个故事，它可能会使用缓存或加载块。由于训练数据集太小，所以在后续版本中我会移除这个机制。

from transformers import AutoModelForCausalLM, AutoTokenizer

How to use in openmind

import torch.nn.functional as F
from torch import Tensor
from openmind import AutoTokenizer, AutoModelForCausalLM, is_torch_npu_available

if is_torch_npu_available():
    device = "npu:0"
else:
    device = "cpu"
tokenizer = AutoTokenizer.from_pretrained("models/TinyLLama-v0")
model = AutoModelForCausalLM.from_pretrained("models/TinyLLama-v0")
model = model.to(device)
prompt = "Give me a short introduction to large language model."
inputs = tokenizer(prompt, return_tensors="pt", return_token_type_ids=False).to(device)
out = model.generate(**inputs, max_new_tokens=80).ravel()
out = tokenizer.decode(out)
print(out)