为何run_t5_mlm_flax.py未生成模型权重等文件？

背景 #

我在复现Hugging Face的T5类掩码语言模型（T5-like span masked-language-modeling）教程，流程分为两步：先运行自定义的分词与配置脚本，再执行官方训练脚本，但训练完成后输出目录的文件远少于预期。

分词与配置脚本（tokenizing_and_configing.py） #

import datasets

from t5_tokenizer_model import SentencePieceUnigramTokenizer
from transformers import T5Config


vocab_size = 32_000
input_sentence_size = None

# 计算数据集总样本数
total_samples = datasets.load_dataset(
    "nthngdy/oscar-mini", name="unshuffled_deduplicated_no", split="train"
).num_rows

# 取总样本的1/30
subset_samples = total_samples // 30

# 加载子集数据
dataset = datasets.load_dataset(
    "nthngdy/oscar-mini",
    name="unshuffled_deduplicated_no",
    split=f"train[:{subset_samples}]",
)

tokenizer = SentencePieceUnigramTokenizer(
    unk_token="<unk>", eos_token="</s>", pad_token="<pad>"
)


# 构建数据集迭代器
def batch_iterator(input_sentence_size=None):
    if input_sentence_size is None:
        input_sentence_size = len(dataset)
    batch_length = 100
    for i in range(0, input_sentence_size, batch_length):
        yield dataset[i : i + batch_length]["text"]


print("Train Tokenizer")
# 训练分词器
tokenizer.train_from_iterator(
    iterator=batch_iterator(input_sentence_size=input_sentence_size),
    vocab_size=vocab_size,
    show_progress=True,
)

# 保存分词器
tokenizer.save("./models/norwegian-t5-base/tokenizer.json")

print("DONE TOKENIZING ")

# 配置模型
config = T5Config.from_pretrained(
    "google/t5-v1_1-small",
    vocab_size=tokenizer.get_vocab_size()
    # "google/t5-v1_1-base", vocab_size=tokenizer.get_vocab_size()
)
config.save_pretrained("./models/norwegian-t5-base")

print("DONE SAVING TOKENIZER ")

依赖的t5_tokenizer_model.py为Hugging Face官方示例中的对应文件。

训练命令 #

完成上述脚本后，执行以下训练命令：

python run_t5_mlm_flax.py \
    --output_dir="./models/norwegian-t5-base" \
    --model_type="t5" \
    --config_name="./models/norwegian-t5-base" \
    --tokenizer_name="./models/norwegian-t5-base" \
    --dataset_name="nthngdy/oscar-mini" \
    --dataset_config_name="unshuffled_deduplicated_no" \
    --max_seq_length="512" \
    --per_device_train_batch_size="32" \
    --per_device_eval_batch_size="32" \
    --adafactor \
    --learning_rate="0.005" \
    --weight_decay="0.001" \
    --warmup_steps="2000" \
    --overwrite_output_dir \
    --logging_steps="500" \
    --save_steps="10000" \
    --eval_steps="2500" \
    --do_train \
    --do_eval

run_t5_mlm_flax.py为Hugging Face官方提供的训练脚本。

当前现象 #

训练过程无报错、流畅完成，但./models/norwegian-t5-base目录下仅生成以下文件：

.
└── norwegian-t5-base
    ├── config.json
    ├── events.out.tfevents.1680920382.ip-172-31-30-81.71782.0.v2
    ├── tokenizer.json
    └── eval_results.json

预期生成的文件 #

我原本期望得到以下文件：

• flax_model.msgpack：微调后的Flax模型权重文件
• tokenizer_config.json：分词器配置文件，包含词汇量大小、特殊token等信息
• training_args.bin：训练参数文件，包含学习率、批次大小等训练时的配置
• merges.txt：分词器子词合并规则文件
• vocab.json：分词器词汇映射关系文件
• train.log：训练过程日志，包含损失值、学习率等指标
• 检查点文件：启用训练检查点后，特定训练步骤的模型权重文件

环境信息 #

使用Amazon AWS p3.2xlarge实例，CUDA版本为cuda_11.2.r11.2/compiler.29618528_0。

内容的提问来源于stack exchange，提问作者littleworth