在兩張 RTX 5090 上,為了極大化訓練效率,我們需要使用 Data Packing(數據打包) 技術。這能確保每個訓練樣本都填滿 4096 或 8192 的長度,不浪費任何顯存空間。
1. 結構化數據的處理策略:分塊 (Chunking)
由於您的數據是 .txt,建議在原始文本中確保每個邏輯塊之間有一個明確的分隔符號(例如 \n\n 或 <|endoftext|>)。
2. 高效能 Packing Data Collator (Python 代碼)
這段代碼會將您的巨大 .txt 讀入,進行分詞,並將多條結構化數據「打包」在一起。這樣可以讓兩張 5090 的計算效率提升 2 倍以上。
</>Python
import torch
from itertools import chain
from datasets import load_dataset
from transformers import AutoTokenizer
# 1. 載入數據與分詞器
tokenizer = AutoTokenizer.from_pretrained("your-llama-7b-path")
tokenizer.pad_token = tokenizer.eos_token
def pack_data(file_path, block_size=4096):
# 使用 streaming 模式處理巨大的 .txt 以節省 RAM
dataset = load_dataset("text", data_files={"train": file_path}, streaming=True)
def tokenize_function(examples):
return tokenizer(examples["text"])
# 分詞
tokenized_dataset = dataset.map(
tokenize_function,
batched=True,
remove_columns=["text"]
)
# 打包邏輯:將 token 序列串接後,按 block_size 切割
def group_texts(examples):
# 串接所有 token
concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
total_length = len(concatenated_examples[list(examples.keys())[0]])
# 捨棄不足一塊的尾巴
total_length = (total_length // block_size) * block_size
# 切割成 block_size 大小的塊
result = {
k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
for k, t in concatenated_examples.items()
}
# 在結構化數據中,labels 就是 input_ids 的複製品 (Causal LM)
result["labels"] = result["input_ids"].copy()
return result
packed_dataset = tokenized_dataset.map(
group_texts,
batched=True,
)
return packed_dataset
# 使用範例
# packed_train_dataset = pack_data("big_data.txt", block_size=4096)
3. 針對 2x 5090 的執行參數優化
處理結構化數據時,請在您的 deepspeed_launcher.sh 中加入以下關鍵微調:
- --gradient_checkpointing True: 雖然 5090 有 32GB 顯存,但處理 4k/8k 長度的全參數微調時,開啟此項可以換取更大的 Batch Size,對結構化數據的穩定收斂很有幫助。
- --weight_decay 0.1: 結構化數據(如代碼)具有高度規律性,高權重衰減能防止模型「死背」特定的變數名稱,轉而學習結構邏輯。
- --lr_scheduler_type linear: 對於增量預訓練,線性衰減(Linear Decay)有時比 Cosine 更能穩定地處理結構化文本。
4. 監控 5090 的關鍵指標
在訓練過程中,除了看 Loss 下降外,請留意以下兩點:
- TFLOPS (每秒浮點運算次數): 5090 應該能跑出極高的數值。如果數值太低,說明數據讀取(CPU/Disk)太慢。
- Memory Usage: 確保兩張卡的顯存佔用均勻。如果一張卡爆掉另一張很空,請檢查
DeepSpeed Stage 3的allgather_bucket_size設定。
5. 接下來的建議
既然數據是結構化的,訓練完畢後,您最需要的是測試模型的「格式正確率」。


















