自宅のRTX3060で小さなLLMを自作してみた

raw_dataset = load_dataset("globis-university/aozorabunko-clean", split="train")
filtered_dataset = raw_dataset.filter(lambda row: row["meta"]["文字遣い種別"] == "新字新仮名")

print(filtered_dataset)
print(f"Filtered dataset size: {filtered_dataset.num_rows:,} entries")

Dataset({
    features: ['text', 'footnote', 'meta'],
    num_rows: 10246
})
Filtered dataset size: 10,246 entries

spm.SentencePieceTrainer.train(
    sentence_iterator=corpus_iterator(filtered_dataset),

    model_prefix=str(OUTPUT_DIR / "aozora_spm_gpt2"),
    model_type="unigram",
    vocab_size=32000,
    character_coverage=0.9995,
    normalization_rule_name="nfkc",
    byte_fallback=True,

    unk_id=0, bos_id=-1, eos_id=-1, pad_id=-1,
    user_defined_symbols=["<|endoftext|>"],

    input_sentence_size=2_000_000,
    shuffle_input_sentence=True,
    train_extremely_large_corpus=True,
    num_threads=8,

    add_dummy_prefix=False,
    remove_extra_whitespaces=False,
    hard_vocab_limit=False,
)

--- Sample 1 ---
Original: 深いおどろきにうたれて、
Token IDs: [874, 11406, 264, 366, 15219, 258]
Pieces: ['深い', 'おどろき', 'に', 'う', 'たれて', '、']
--- Sample 2 ---
Original: 名高いウェストミンスターに
Token IDs: [15899, 14107, 2047, 10438, 7545, 264]
Pieces: ['名高い', 'ウェ', 'スト', 'ミン', 'スター', 'に']
--- Sample 3 ---
Original: 真鍮や石の記念碑となって
Token IDs: [21674, 278, 3600, 5815, 9470, 1070]
Pieces: ['真鍮', 'や', '石の', '記念', '碑', 'となって']
--- Sample 4 ---
Original: すべての王侯貴族が集まっているのをみれば、
Token IDs: [2200, 965, 6015, 5559, 262, 26889, 260, 266, 5117, 258]
Pieces: ['すべての', '王', '侯', '貴族', 'が', '集まっている', 'の', 'を', 'みれば', '、']
--- Sample 5 ---
Original: 今はさげすみも、ほこりも、見栄もない。
Token IDs: [1291, 12277, 3254, 268, 258, 18540, 268, 258, 340, 2529, 510, 259]
Pieces: ['今は', 'さげ', 'すみ', 'も', '、', 'ほこり', 'も', '、', '見', '栄', 'もない', '。']

config = GPT2Config(
    vocab_size=vocab_size,
    n_positions=block_size,
    n_ctx=block_size,
    n_embd=512,
    n_layer=8,
    n_head=8,
    bos_token_id=endoftext_id,
    eos_token_id=endoftext_id,
    pad_token_id=endoftext_id,
)

model = GPT2LMHeadModel(config)
model.config.use_cache = False

tokenized = dataset_dict.map(
    encode_with_sentencepiece,
    batched=True,
    remove_columns=original_columns,
    desc="Tokenizing with SentencePiece",
)

lm_datasets = tokenized.map(
    group_texts,
    batched=True,
    desc="Grouping texts into fixed-size blocks",
)

training_args = TrainingArguments(
    output_dir=str(output_dir),
    num_train_epochs=1,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=16,
    learning_rate=5e-4,
    warmup_steps=500,
    weight_decay=0.1,
    logging_steps=100,
    eval_steps=500,
    save_steps=500,
    evaluation_strategy="steps",
    save_strategy="steps",
    fp16=torch.cuda.is_available(),
    report_to=[],
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=lm_datasets["train"],
    eval_dataset=lm_datasets["test"],
    data_collator=default_data_collator,
)