96GB 的原始文本在 Tokenize 之後，轉換為 Uint16 的二進制數據大約會佔用 150GB - 200GB 的磁碟空間。在兩張 RTX 5090 上運行 7B 模型時，使用 Memory Mapping (mmap) 是確保 GPU 不會因為等待 I/O 而閒置的最強方案。
以下是一個簡化但工業級的實作方案，分為 「轉換腳本」 與 「加載器」 兩部分。
第一步：將 96GB 文本轉換為二進制 (.bin)
我們使用numpy.memmap 來創建一個巨大的二進制文件，並配合 indices 文件記錄數據位置。</>Python
import numpy as np
from transformers import AutoTokenizer
import os
def preprocess_to_bin(text_files, output_bin, tokenizer_path):
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
# 使用 uint16 因為多數詞表大小 < 65535
# 如果詞表 > 65k，請改用 np.uint32
with open(output_bin, 'wb') as f:
for file_path in text_files:
with open(file_path, 'r', encoding='utf-8') as tf:
for line in tf:
tokens = tokenizer.encode(line.strip())
if tokens:
# 將 token 寫入二進制文件
np.array(tokens, dtype=np.uint16).tofile(f)
# 執行轉換 (這可能需要數小時，視 CPU 速度而定)
# preprocess_to_bin(["data1.txt", "data2.txt"], "train_data.bin", "llama-3-tokenizer")
第二步：高效的 mmap 數據加載器
這段代碼會模擬 Dataset 類，讓 PyTorch 能夠以「零內存佔用」的方式隨機存取 200GB 的數據。
</>Python
import torch
from torch.utils.data import Dataset
import numpy as np
class MmapDataset(Dataset):
def __init__(self, bin_path, block_size=2048):
self.bin_path = bin_path
self.block_size = block_size
# 使用 memmap 模式讀取，這不會把文件載入 RAM
# offset=0, mode='r' 代表唯讀映射
self.data = np.memmap(bin_path, dtype=np.uint16, mode='r')
# 計算總共有多少個訓練樣本 (Total Tokens / Block Size)
self.num_samples = len(self.data) // block_size
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
# 根據索引計算起始位置
start = idx * self.block_size
end = start + self.block_size
# 從硬碟映射中切片 (Slice)，操作系統會自動處理頁面快取
chunk = self.data[start:end].astype(np.int64)
# 轉換為 PyTorch Tensor
x = torch.from_numpy(chunk)
# 在預訓練中，Label 通常就是 Input 右移一位 (Causal Language Modeling)
# 這裡簡化處理，直接回傳同一段，在訓練循環中再處理 shift
return x
# 使用方法
dataset = MmapDataset("train_data.bin", block_size=2048)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=4,
shuffle=True,
num_workers=4,
pin_memory=True # 必開，加速數據傳輸到 5090
)
針對 2x 5090 的專業調優建議
- 數據類型選擇： 如果你的 Tokenizer 是 Llama-3 (詞表 128k)，
np.uint16會溢出，請務必改用 np.uint32。雖然這會讓數據集體積增加一倍，但能保證數據正確性。 - 多卡分布式讀取： 在使用
DistributedDataParallel或DeepSpeed時，記得在DataLoader中加入DistributedSampler：
</>Python
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler, ...)
- 磁碟性能： 96GB 文本轉成 bin 後很大，請確保你的 SSD 剩餘空間大於 300GB。對於 5090 這種等級的卡，建議使用 NVMe Gen4 或 Gen5 SSD，否則
shuffle=True時的隨機讀取可能會導致卡頓。
💡 小提醒：24 小時不間斷訓練
兩張 5090 滿載預訓練 7B 模型是非常吃電的。建議你在啟動腳本前，先用 nvidia-smi 限制每張卡的功耗（例如限制在 450W），這能防止突發大電流導致跳電，且對訓練速度的影響通常不到 5%。
