注意
转到页面底部 下载完整示例代码。
使用 PyTorch 和 Ray Train 进行大规模分布式训练#
作者: Ricardo Decal
本教程展示了如何使用 Ray Train 和 Ray Data 将 PyTorch 训练任务分布到多个 GPU 上，以实现可扩展且适用于生产环境的模型训练。
你将学习如何：
使用 PyTorch 和 Hugging Face Transformers 预训练一个 GPT-2（约 1.24 亿参数）语言模型。
通过 Ray Train 以最少的代码更改将训练任务分布到多个 GPU 上。
使用 Ray Data 的分布式工作节点从 Hugging Face 数据集流式传输训练数据。
保存和加载分布式检查点（Checkpoints）。
通过微小的代码调整，实现从单节点到多节点集群的扩展。
利用异构集群优化成本和性能。
使用 Ray 仪表盘监控训练过程。
先决条件
PyTorch v2.9+。
Ray Train (
ray[train]) v2.52.1+。tiktoken,datasets和transformers(Hugging Face)。建议使用一个或多个 GPU，但非必需。本教程在
g4dn.12xlarge实例上进行了测试，该实例配备 4 个 NVIDIA T4 GPU（每个 GPU 16GB 显存）。
Ray Train 是一个用于分布式深度学习的可扩展框架。Ray Train 构建于 Ray 之上，Ray 是一个统一的 AI 和 Python 应用扩展框架，简化了分布式计算的复杂性。Ray 也是开源的，并且是 PyTorch 基金会的一部分。
Ray Train 使你无需重写训练循环即可从单个 GPU 扩展到数百个 GPU。结合用于流式数据摄取的 Ray Data，你可以获得一个端到端的分布式训练流水线，该流水线处理数据加载、分片（sharding）、梯度同步、检查点保存和容错机制。
设置#
要安装依赖项，请运行 pip install "ray[train]" torch tiktoken datasets transformers。
然后，导入所需的库
import os
import tempfile
import time
import numpy as np
import ray
import ray.train
import tiktoken
import torch
from datasets import load_dataset
from ray.train import CheckpointConfig, RunConfig, ScalingConfig
from ray.train.torch import TorchTrainer
from transformers import GPT2Config, GPT2LMHeadModel
# Enable smoke test to run this tutorial quickly.
SMOKE_TEST = True
# Reduce Ray Data verbosity
ray.data.DataContext.get_current().enable_progress_bars = False
ray.data.DataContext.get_current().print_on_execution_start = False
使用 Ray Data 加载数据集#
本教程使用 Wikitext-103 数据集，该数据集收集了维基百科中经过验证的优质和特色文章，包含超过 1 亿个 token。
ray.data.from_huggingface() 函数将 Hugging Face 数据集转换为 Ray Dataset，从而实现在所有可用节点上的分布式流式传输和预处理。
hf_ds = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1")
train_ds = ray.data.from_huggingface(hf_ds["train"])
val_ds = ray.data.from_huggingface(hf_ds["validation"])
# Limit dataset size for fast iteration during smoke tests.
if SMOKE_TEST:
train_ds = train_ds.limit(2500)
val_ds = val_ds.limit(2500)
print(f"Dataset schema:\n{train_ds.schema()}")
Downloading readme: 0.00B [00:00, ?B/s]
Downloading readme: 10.5kB [00:00, 33.4MB/s]
Downloading data: 0%| | 0.00/733k [00:00<?, ?B/s]
Downloading data: 100%|██████████| 733k/733k [00:00<00:00, 3.27MB/s]
Downloading data: 100%|██████████| 733k/733k [00:00<00:00, 3.24MB/s]
Downloading data: 0%| | 0.00/157M [00:00<?, ?B/s]
Downloading data: 7%|▋ | 10.5M/157M [00:00<00:05, 28.2MB/s]
Downloading data: 20%|██ | 31.5M/157M [00:00<00:01, 71.0MB/s]
Downloading data: 33%|███▎ | 52.4M/157M [00:11<00:29, 3.54MB/s]
Downloading data: 40%|████ | 62.9M/157M [00:20<00:40, 2.34MB/s]
Downloading data: 47%|████▋ | 73.4M/157M [00:21<00:28, 2.89MB/s]
Downloading data: 60%|██████ | 94.4M/157M [00:22<00:12, 5.01MB/s]
Downloading data: 73%|███████▎ | 115M/157M [00:30<00:11, 3.56MB/s]
Downloading data: 80%|████████ | 126M/157M [00:38<00:12, 2.58MB/s]
Downloading data: 87%|████████▋ | 136M/157M [00:43<00:08, 2.54MB/s]
Downloading data: 100%|██████████| 157M/157M [00:43<00:00, 3.63MB/s]
Downloading data: 0%| | 0.00/157M [00:00<?, ?B/s]
Downloading data: 7%|▋ | 10.5M/157M [00:00<00:08, 18.1MB/s]
Downloading data: 20%|██ | 31.5M/157M [00:00<00:02, 53.7MB/s]
Downloading data: 33%|███▎ | 52.4M/157M [00:00<00:01, 78.9MB/s]
Downloading data: 47%|████▋ | 73.4M/157M [00:00<00:00, 98.6MB/s]
Downloading data: 60%|██████ | 94.4M/157M [00:01<00:00, 117MB/s]
Downloading data: 73%|███████▎ | 115M/157M [00:01<00:00, 124MB/s]
Downloading data: 87%|████████▋ | 136M/157M [00:01<00:00, 122MB/s]
Downloading data: 100%|██████████| 157M/157M [00:01<00:00, 125MB/s]
Downloading data: 100%|██████████| 157M/157M [00:01<00:00, 98.0MB/s]
Downloading data: 0%| | 0.00/657k [00:00<?, ?B/s]
Downloading data: 100%|██████████| 657k/657k [00:00<00:00, 2.90MB/s]
Downloading data: 100%|██████████| 657k/657k [00:00<00:00, 2.88MB/s]
Generating test split: 0%| | 0/4358 [00:00<?, ? examples/s]
Generating test split: 100%|██████████| 4358/4358 [00:00<00:00, 519829.84 examples/s]
Generating train split: 0%| | 0/1801350 [00:00<?, ? examples/s]
Generating train split: 5%|▍ | 89000/1801350 [00:00<00:01, 876067.07 examples/s]
Generating train split: 10%|▉ | 180000/1801350 [00:00<00:01, 890714.32 examples/s]
Generating train split: 15%|█▌ | 272000/1801350 [00:00<00:01, 901170.58 examples/s]
Generating train split: 20%|██ | 364000/1801350 [00:00<00:01, 903955.86 examples/s]
Generating train split: 25%|██▌ | 455000/1801350 [00:00<00:01, 901383.30 examples/s]
Generating train split: 30%|███ | 547000/1801350 [00:00<00:01, 904094.12 examples/s]
Generating train split: 36%|███▌ | 640000/1801350 [00:00<00:01, 908647.34 examples/s]
Generating train split: 41%|████ | 732000/1801350 [00:00<00:01, 907797.18 examples/s]
Generating train split: 46%|████▌ | 824000/1801350 [00:00<00:01, 907669.33 examples/s]
Generating train split: 53%|█████▎ | 957675/1801350 [00:01<00:00, 896733.36 examples/s]
Generating train split: 58%|█████▊ | 1051675/1801350 [00:01<00:00, 904160.74 examples/s]
Generating train split: 63%|██████▎ | 1143675/1801350 [00:01<00:00, 903189.81 examples/s]
Generating train split: 69%|██████▊ | 1235675/1801350 [00:01<00:00, 905464.30 examples/s]
Generating train split: 74%|███████▍ | 1328675/1801350 [00:01<00:00, 909461.94 examples/s]
Generating train split: 81%|████████▏ | 1465675/1801350 [00:01<00:00, 908811.34 examples/s]
Generating train split: 87%|████████▋ | 1558675/1801350 [00:01<00:00, 911169.46 examples/s]
Generating train split: 94%|█████████▍| 1695675/1801350 [00:01<00:00, 909874.12 examples/s]
Generating train split: 99%|█████████▉| 1788675/1801350 [00:01<00:00, 912203.98 examples/s]
Generating train split: 100%|██████████| 1801350/1801350 [00:01<00:00, 905866.40 examples/s]
Generating validation split: 0%| | 0/3760 [00:00<?, ? examples/s]
Generating validation split: 100%|██████████| 3760/3760 [00:00<00:00, 729950.62 examples/s]
Downloading readme: 0.00B [00:00, ?B/s]
Downloading readme: 10.5kB [00:00, 11.2MB/s]
2026-07-08 22:20:16,399 WARNING services.py:2213 -- WARNING: The object store is using /tmp/ray instead of /dev/shm because /dev/shm has only 2147467264 bytes available. This will harm performance! You may be able to free up space by deleting files in /dev/shm. If you are inside a Docker container, you can increase /dev/shm size by passing '--shm-size=10.24gb' to 'docker run' (or add it to the run_options list in a Ray cluster config). Make sure to set this to more than 30% of available RAM.
2026-07-08 22:20:18,596 INFO worker.py:2003 -- Started a local Ray instance. View the dashboard at 127.0.0.1:8265
/usr/local/lib/python3.10/dist-packages/ray/_private/worker.py:2051: FutureWarning: Tip: In future versions of Ray, Ray will no longer override accelerator visible devices env var if num_gpus=0 or num_gpus=None (default). To enable this behavior and turn off this error message, set RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0
warnings.warn(
Dataset schema:
Column Type
------ ----
text string
架构看起来可能像这样
Column Type
------ ----
text string
这意味着该数据集有一列名为 text，且为字符串类型。
检查原始数据#
使用 take(n) 获取少量行进行检查。每一行都是一个字典，以列名为键。
print("--- Raw data sample ---")
sample = train_ds.take(2)
for i, row in enumerate(sample):
text_preview = (row["text"][:120] + "...") if len(row["text"]) > 120 else row["text"]
print(f" Row {i}: {text_preview!r}")
--- Raw data sample ---
2026-07-08 22:20:21,877 INFO dataset.py:3818 -- Tip: Use `take_batch()` instead of `take() / show()` to return records in pandas or numpy batch format.
2026-07-08 22:20:21,885 INFO logging.py:416 -- Registered dataset logger for dataset dataset_4_0
2026-07-08 22:20:21,900 WARNING resource_manager.py:169 -- ⚠️ Ray's object store is configured to use only 5.3% of available memory (9.3GiB out of 176.5GiB total). For optimal Ray Data performance, we recommend setting the object store to at least 50% of available memory. You can do this by setting the 'object_store_memory' parameter when calling ray.init() or by setting the RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION environment variable.
2026-07-08 22:20:21,900 WARNING __init__.py:28 -- Progress bars disabled. To enable, set `ray.data.DataContext.get_current().enable_progress_bars = True`.
2026-07-08 22:20:22,595 INFO streaming_executor.py:294 -- ✔️ Dataset dataset_4_0 execution finished in 0.70 seconds
Row 0: ''
Row 1: ' = Valkyria Chronicles III = \n'
你将看到类似这样的输出
Row 0: ''
Row 1: ' = Valkyria Chronicles III = '
Wikitext-103 中的每一行都是维基百科文章的一行。连续的行属于同一篇文章，空行用于分隔段落。新文章以 = Article Title = 这样的标题行开始。下方的分词步骤会在每个标题行之前插入一个 <|endoftext|> 分隔符 token，以便模型学习在文章边界处重置上下文。
对数据进行分词和分块#
语言模型处理的是固定长度的 token ID 序列。预处理步骤将原始文本转换为用于下一 token 预测的 token ID 序列。
本教程使用带有 GPT-2 编码（词汇表大小 50,257）的 tiktoken。tiktoken 是一个快速、独立的 tokenizer，不依赖 Hugging Face 的 transformers 库。
tokenize_and_chunk 函数执行以下操作：
对每一批文本进行分词，并连接成单个流。文章标题行（例如
= Article Title =）会触发一个<|endoftext|>分隔符，以便模型在文章边界处重置上下文。将流拆分为固定长度为
block_size个 token 的块。返回每个块的
input_ids。在训练期间，同一个张量同时用作输入和标签，因为GPT2LMHeadModel在计算交叉熵损失时会在内部平移标签。
BLOCK_SIZE = 256
VOCAB_SIZE = 50257
encoding = tiktoken.get_encoding("gpt2")
EOT_TOKEN = encoding.eot_token # <|endoftext|> token ID (50256)
def _is_article_title(text: str) -> bool:
"""Detect Wikitext article title lines like ' = Some Title = '."""
stripped = text.strip()
return stripped.startswith("= ") and stripped.endswith(" =") and not stripped.startswith("= =")
def tokenize_and_chunk(batch: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
"""Tokenize text and split into fixed-length chunks for language modeling."""
# Reconstruct the original text stream by joining rows with newlines.
# Article title lines signal new articles, so we insert an
# <|endoftext|> separator before them.
all_tokens: list[int] = []
for text in batch["text"]:
if _is_article_title(text):
all_tokens.append(EOT_TOKEN)
all_tokens.extend(encoding.encode_ordinary(text + "\n"))
# Split into fixed-length chunks of block_size tokens.
num_chunks = len(all_tokens) // BLOCK_SIZE
all_tokens = all_tokens[: num_chunks * BLOCK_SIZE]
if num_chunks == 0:
return {"input_ids": []}
tokens_array = np.array(all_tokens, dtype=np.int64).reshape(num_chunks, BLOCK_SIZE)
return {"input_ids": tokens_array}
使用 map_batches() 应用分词。此操作是惰性（lazy）的，意味着 Ray Data 会推迟执行，直到下游消费者请求结果。惰性执行使 Ray 能够在任何工作开始前优化整个流水线。
# These do not trigger execution.
train_ds = train_ds.map_batches(tokenize_and_chunk, batch_format="numpy")
val_ds = val_ds.map_batches(tokenize_and_chunk, batch_format="numpy")
使用 take(2) 检查分词后的输出
print("--- After tokenization ---")
tokenized_sample = train_ds.take(2)
for i, row in enumerate(tokenized_sample):
ids = row["input_ids"]
print(f" Row {i}: input_ids shape={ids.shape}, first 10 tokens={ids[:10].tolist()}")
print(f" Decoded: {encoding.decode(ids[:30].tolist())!r}...")
--- After tokenization ---
2026-07-08 22:20:23,408 INFO logging.py:416 -- Registered dataset logger for dataset dataset_7_0
2026-07-08 22:20:23,929 INFO streaming_executor.py:294 -- ✔️ Dataset dataset_7_0 execution finished in 0.52 seconds
Row 0: input_ids shape=(256,), first 10 tokens=[198, 50256, 796, 569, 18354, 7496, 17740, 6711, 796, 220]
Decoded: '\n<|endoftext|> = Valkyria Chronicles III = \n\n\n Senjō no Valkyria 3 : Unrecorded Chronicles ( Japanese : 戦'...
Row 1: input_ids shape=(256,), first 10 tokens=[33687, 5303, 18024, 6909, 764, 317, 1588, 1074, 286, 8786]
Decoded: " Takeshi Ozawa . A large team of writers handled the script . The game 's opening theme was sung by May 'n . \n\n It"...
每一行现在都包含一个长度为 256 个 token 的固定长度 input_ids 数组。
流式执行#
在内部，Ray 将数据划分为块（blocks）并将其分发给工作节点。这种基于块的架构实现了流式执行：一旦某个阶段输出了一个块，下一个阶段就可以立即开始处理它，而无需等待前面的阶段完成整个数据集的处理。这意味着上述 map_batches 分词操作与训练循环在流式流水线中并行运行，因此整个数据集无需一次性全部加载到内存中。
当训练开始时，Ray Data 会记录执行计划。对于本教程，一种可能的计划是：
Execution plan: InputDataBuffer[Input]
-> TaskPoolMapOperator[MapBatches(tokenize_and_chunk)]
-> OutputSplitter[split(4, equal=True)]
这清楚地告诉你 Ray Data 将如何通过分词进行流式传输，并将数据分发给 4 个训练工作节点。
定义 Transformer 模型#
该模型是一个仅解码器（decoder-only）的 Transformer 语言模型，使用了 Hugging Face 的 GPT2LMHeadModel。下面的超参数适用于标准的 GPT-2 “小”架构。
def create_model():
"""Create a GPT-2 small model with random weights."""
model = GPT2LMHeadModel(GPT2Config(
vocab_size=VOCAB_SIZE,
n_positions=BLOCK_SIZE,
n_embd=768,
n_layer=12,
n_head=12,
))
model.loss_type = "ForCausalLM"
return model
验证模型大小
model = create_model()
num_params = sum(p.numel() for p in model.parameters())
print(f"Model parameters: {num_params / 1e6:.1f}M")
del model # Free memory before training
Model parameters: 123.8M
你可以看到大约有 1.238 亿个参数。
定义分布式训练函数#
训练函数在每个工作进程上运行。Ray Train 管理分布式设置：它将模型封装在 DistributedDataParallel 中，在工作节点间分片数据，并自动同步梯度。
Ray Train 的关键集成点是：
ray.train.get_dataset_shard("train")：获取该工作节点的特定数据集分片，Ray Data 会自动将数据集在所有工作节点间拆分。ray.train.torch.prepare_model(model)：将模型封装在DistributedDataParallel中，并将其移动到正确的 GPU 上。shard.iter_torch_batches(batch_size=...)：返回一个dict[str, torch.Tensor]批次的迭代器，张量会自动放置在工作节点的 GPU 上。设置prefetch_batches=2可以预先抓取 2 个批次。ray.train.report(metrics, checkpoint=...)：向驱动程序（driver）报告指标并保存检查点。
def train_func_per_worker(config: dict):
"""Training function executed by each distributed worker."""
lr = config["lr"]
weight_decay = config["weight_decay"]
max_grad_norm = config["max_grad_norm"]
epochs = config["epochs"]
batch_size = config["batch_size_per_worker"]
max_steps_per_epoch = config.get("max_steps_per_epoch")
# --- Data -----------------------------------------------------------
# Each worker gets an automatic shard of the dataset.
train_data_shard = ray.train.get_dataset_shard("train")
val_data_shard = ray.train.get_dataset_shard("validation")
# --- Model ----------------------------------------------------------
model = create_model()
# prepare_model wraps the model in DistributedDataParallel and places
# it on the correct device.
model = ray.train.torch.prepare_model(model)
optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
# --- Training loop --------------------------------------------------
for epoch in range(epochs):
model.train()
train_loss_sum = 0.0
train_batches = 0
train_tokens = 0
epoch_start = time.perf_counter()
# iter_torch_batches returns dicts of tensors already on the GPU.
for batch in train_data_shard.iter_torch_batches(
batch_size=batch_size, dtypes=torch.long, prefetch_batches=2
):
input_ids = batch["input_ids"]
# GPT2LMHeadModel shifts labels internally to align each
# position with the next token, so we can use input_ids as
# both the input and the labels.
out = model(input_ids=input_ids, labels=input_ids)
loss = out.loss
optimizer.zero_grad()
loss.backward()
# Gradient clipping for training stability
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)
optimizer.step()
train_loss_sum += loss.item()
train_batches += 1
train_tokens += input_ids.numel()
if max_steps_per_epoch and train_batches >= max_steps_per_epoch:
break
train_elapsed = time.perf_counter() - epoch_start
avg_train_loss = train_loss_sum / max(train_batches, 1)
# --- Validation -----------------------------------------------------
model.eval()
val_loss_sum = 0.0
val_batches = 0
with torch.no_grad():
for batch in val_data_shard.iter_torch_batches(
batch_size=batch_size, dtypes=torch.long, prefetch_batches=2
):
input_ids = batch["input_ids"]
out = model(input_ids=input_ids, labels=input_ids)
loss = out.loss
val_loss_sum += loss.item()
val_batches += 1
if max_steps_per_epoch and val_batches >= max_steps_per_epoch:
break
avg_val_loss = val_loss_sum / max(val_batches, 1)
epoch_elapsed = time.perf_counter() - epoch_start
# --- Report metrics and save checkpoint ------------------------------
metrics = {
"train_loss": round(avg_train_loss, 4),
"val_loss": round(avg_val_loss, 4),
"epoch": epoch,
"epoch_time_sec": round(epoch_elapsed, 2),
"epoch_tokens": train_tokens,
"tokens_per_sec": round(train_tokens / max(train_elapsed, 1e-6), 2),
}
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
torch.save(
{
"epoch": epoch,
"model_state_dict": model.module.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
os.path.join(temp_checkpoint_dir, "checkpoint.pt"),
)
checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
ray.train.report(metrics=metrics, checkpoint=checkpoint)
配置并启动分布式训练#
TorchTrainer 将一切整合在一起。运行 trainer.fit() 最终会触发完整数据流水线和训练循环的执行。Trainer 接受：
train_func_per_worker：每个工作节点执行的函数。train_loop_config：转发给训练函数的超参数字典。datasets：Ray Datasets 字典。Ray Train 会自动在工作节点间拆分每个数据集。scaling_config：指定工作节点的数量以及是否使用 GPU。
设置 num_workers=4 会启动 4 个并行工作节点，每个 GPU 一个。Ray Train 在后台处理 torch.distributed 初始化、NCCL 后端设置和 DistributedDataParallel 封装。在日志中，你会看到每个工作节点被分配了一个排名（rank）和设备。
Started training worker group of size 4:
* (ip=10.0.176.183, pid=25636) world_rank=0, local_rank=0, node_rank=0
* (ip=10.0.176.183, pid=25637) world_rank=1, local_rank=1, node_rank=0
...
Moving model to device: cuda:0
Wrapping provided model in DistributedDataParallel.
batch_size_per_worker 是每个工作节点在每个梯度步骤处理的序列数。使用 4 个工作节点，每个工作节点的批大小为 16，有效的全局批大小为 4 × 16 = 64 个序列，即每个优化步骤 64 × 256 = 4,096 个 token。
USE_GPU = torch.cuda.is_available()
NUM_WORKERS = max(torch.cuda.device_count(), 1) # One worker per available GPU
NUM_EPOCHS = 5
BATCH_SIZE_PER_WORKER = 16
LR = 3e-4
WEIGHT_DECAY = 0.1
MAX_GRAD_NORM = 1.0
trainer = TorchTrainer(
train_loop_per_worker=train_func_per_worker,
train_loop_config={
"lr": LR,
"weight_decay": WEIGHT_DECAY,
"max_grad_norm": MAX_GRAD_NORM,
"epochs": NUM_EPOCHS,
"batch_size_per_worker": BATCH_SIZE_PER_WORKER,
"max_steps_per_epoch": 5 if SMOKE_TEST else None,
},
# Register the datasets,
datasets={"train": train_ds, "validation": val_ds},
scaling_config=ScalingConfig(
num_workers=NUM_WORKERS,
use_gpu=USE_GPU,
),
run_config=RunConfig(
name="gpt2-small-pretraining",
storage_path="/tmp/ray-train-checkpoints",
),
)
result = trainer.fit()
(TrainController pid=7974) Requesting resources: {'GPU': 1} * 4
(TrainController pid=7974) Attempting to start training worker group of size 4 with the following resources: [{'GPU': 1}] * 4
(RayTrainWorker pid=8094) Setting up process group for: env:// [rank=0, world_size=4]
(TrainController pid=7974) Started training worker group of size 4:
(TrainController pid=7974) - (ip=172.17.0.2, pid=8094) world_rank=0, local_rank=0, node_rank=0
(TrainController pid=7974) - (ip=172.17.0.2, pid=8095) world_rank=1, local_rank=1, node_rank=0
(TrainController pid=7974) - (ip=172.17.0.2, pid=8096) world_rank=2, local_rank=2, node_rank=0
(TrainController pid=7974) - (ip=172.17.0.2, pid=8097) world_rank=3, local_rank=3, node_rank=0
(RayTrainWorker pid=8094) Moving model to device: cuda:0
(RayTrainWorker pid=8094) Wrapping provided model in DistributedDataParallel.
(SplitCoordinator pid=8558) Registered dataset logger for dataset train_8_0
(SplitCoordinator pid=8558) ⚠️ Ray's object store is configured to use only 5.3% of available memory (9.3GiB out of 176.5GiB total). For optimal Ray Data performance, we recommend setting the object store to at least 50% of available memory. You can do this by setting the 'object_store_memory' parameter when calling ray.init() or by setting the RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION environment variable.
(SplitCoordinator pid=8558) Progress bars disabled. To enable, set `ray.data.DataContext.get_current().enable_progress_bars = True`.
(SplitCoordinator pid=8558) ✔️ Dataset train_8_0 execution finished in 1.66 seconds
(RayTrainWorker pid=8094) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/tmp/ray-train-checkpoints/gpt2-small-pretraining/checkpoint_2026-07-08_22-20-46.310783)
(RayTrainWorker pid=8094) Reporting training result 1: TrainingReport(checkpoint=Checkpoint(filesystem=local, path=/tmp/ray-train-checkpoints/gpt2-small-pretraining/checkpoint_2026-07-08_22-20-46.310783), metrics={'train_loss': 9.8328, 'val_loss': 8.9673, 'epoch': 0, 'epoch_time_sec': 5.11, 'epoch_tokens': 20480, 'tokens_per_sec': 4650.26}, validation=False)
(SplitCoordinator pid=8559) Registered dataset logger for dataset validation_10_0
(SplitCoordinator pid=8559) ⚠️ Ray's object store is configured to use only 5.3% of available memory (9.3GiB out of 176.5GiB total). For optimal Ray Data performance, we recommend setting the object store to at least 50% of available memory. You can do this by setting the 'object_store_memory' parameter when calling ray.init() or by setting the RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION environment variable.
(SplitCoordinator pid=8559) Progress bars disabled. To enable, set `ray.data.DataContext.get_current().enable_progress_bars = True`.
(SplitCoordinator pid=8559) ✔️ Dataset validation_10_0 execution finished in 0.22 seconds
(SplitCoordinator pid=8558) Registered dataset logger for dataset train_8_1
(SplitCoordinator pid=8558) ✔️ Dataset train_8_1 execution finished in 0.19 seconds
(RayTrainWorker pid=8095) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/tmp/ray-train-checkpoints/gpt2-small-pretraining/checkpoint_2026-07-08_22-20-53.304984) [repeated 4x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.rayai.org.cn/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)
(RayTrainWorker pid=8095) Reporting training result 2: TrainingReport(checkpoint=Checkpoint(filesystem=local, path=/tmp/ray-train-checkpoints/gpt2-small-pretraining/checkpoint_2026-07-08_22-20-53.304984), metrics={'train_loss': 8.5157, 'val_loss': 8.2642, 'epoch': 1, 'epoch_time_sec': 3.13, 'epoch_tokens': 20480, 'tokens_per_sec': 8186.97}, validation=False) [repeated 4x across cluster]
(SplitCoordinator pid=8559) Registered dataset logger for dataset validation_10_1
(SplitCoordinator pid=8559) ✔️ Dataset validation_10_1 execution finished in 0.16 seconds
(SplitCoordinator pid=8558) Registered dataset logger for dataset train_8_2
(SplitCoordinator pid=8558) ✔️ Dataset train_8_2 execution finished in 0.20 seconds
(RayTrainWorker pid=8094) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/tmp/ray-train-checkpoints/gpt2-small-pretraining/checkpoint_2026-07-08_22-21-00.368019) [repeated 4x across cluster]
(RayTrainWorker pid=8094) Reporting training result 3: TrainingReport(checkpoint=Checkpoint(filesystem=local, path=/tmp/ray-train-checkpoints/gpt2-small-pretraining/checkpoint_2026-07-08_22-21-00.368019), metrics={'train_loss': 7.6497, 'val_loss': 7.7739, 'epoch': 2, 'epoch_time_sec': 3.13, 'epoch_tokens': 20480, 'tokens_per_sec': 8164.15}, validation=False) [repeated 4x across cluster]
(SplitCoordinator pid=8559) Registered dataset logger for dataset validation_10_2
(SplitCoordinator pid=8559) ✔️ Dataset validation_10_2 execution finished in 0.15 seconds
(SplitCoordinator pid=8558) Registered dataset logger for dataset train_8_3
(SplitCoordinator pid=8558) ✔️ Dataset train_8_3 execution finished in 0.19 seconds
(RayTrainWorker pid=8096) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/tmp/ray-train-checkpoints/gpt2-small-pretraining/checkpoint_2026-07-08_22-21-07.338114) [repeated 4x across cluster]
(RayTrainWorker pid=8096) Reporting training result 4: TrainingReport(checkpoint=Checkpoint(filesystem=local, path=/tmp/ray-train-checkpoints/gpt2-small-pretraining/checkpoint_2026-07-08_22-21-07.338114), metrics={'train_loss': 7.1617, 'val_loss': 7.7314, 'epoch': 3, 'epoch_time_sec': 3.16, 'epoch_tokens': 20480, 'tokens_per_sec': 8068.59}, validation=False) [repeated 4x across cluster]
(SplitCoordinator pid=8559) Registered dataset logger for dataset validation_10_3
(SplitCoordinator pid=8559) ✔️ Dataset validation_10_3 execution finished in 0.16 seconds
(SplitCoordinator pid=8558) Registered dataset logger for dataset train_8_4
(SplitCoordinator pid=8558) ✔️ Dataset train_8_4 execution finished in 0.25 seconds
(RayTrainWorker pid=8094) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/tmp/ray-train-checkpoints/gpt2-small-pretraining/checkpoint_2026-07-08_22-21-14.340814) [repeated 4x across cluster]
(RayTrainWorker pid=8094) Reporting training result 5: TrainingReport(checkpoint=Checkpoint(filesystem=local, path=/tmp/ray-train-checkpoints/gpt2-small-pretraining/checkpoint_2026-07-08_22-21-14.340814), metrics={'train_loss': 7.1693, 'val_loss': 7.7974, 'epoch': 4, 'epoch_time_sec': 3.2, 'epoch_tokens': 20480, 'tokens_per_sec': 7928.22}, validation=False) [repeated 4x across cluster]
(SplitCoordinator pid=8559) Registered dataset logger for dataset validation_10_4
(SplitCoordinator pid=8559) ✔️ Dataset validation_10_4 execution finished in 0.15 seconds
检查结果#
训练结束后，Result 对象包含最终指标和检查点。result.metrics 来自最后一次 ray.train.report() 调用。result.checkpoint 包含来自最后一次 ray.train.report() 调用的检查点。
print("\nTraining finished!")
Training finished!
result.metrics 包含了最后一次 ray.train.report() 调用中的指标字典。
{'train_loss': 7.0646, 'val_loss': 7.6051, 'epoch': 4,
'epoch_time_sec': 12.34, 'epoch_tokens': 20480, 'tokens_per_sec': 1759.8}
每个工作节点的日志显示了每个 epoch 的训练损失、验证损失和吞吐量指标。使用随机权重且只进行了少量步骤的情况下，预计会有较高的损失（约 10-11）。
检查点（Checkpointing）#
在生产环境训练中，你可以启用检查点功能，以使训练任务能够应对意外故障。检查点允许你利用 容错 部分中描述的 Ray Train 容错机制。
Ray Train 提供了多种检查点优化功能。异步上传允许你在检查点在后台流式传输到远程存储时继续训练。分布式检查点并行上传每个工作节点的分片，避免了将数据汇总到单个工作节点内存中，从而降低了大型模型出现 OOM（内存溢出）错误的风险。
有关 Ray Train 检查点功能的完整指南，请参阅 Ray Train 检查点指南。
扩展到多节点集群#
上述代码运行在单台 4 GPU 机器上。扩展到多节点集群仅需进行两项更改：
增加 ``num_workers`` 以匹配集群中 GPU 的总数。
设置共享存储路径，以便所有节点都能访问检查点。
例如，要在 4 个节点组成的集群上进行训练，每个节点 4 个 GPU（总共 16 个 GPU）：
trainer = TorchTrainer(
train_loop_per_worker=train_func_per_worker,
train_loop_config={...},
datasets={"train": train_ds, "validation": val_ds},
scaling_config=ScalingConfig(
num_workers=16, # 4 nodes x 4 GPUs
use_gpu=True,
),
run_config=RunConfig(
# Shared storage accessible from all nodes
storage_path="s3://my-bucket/ray-checkpoints",
checkpoint_config=CheckpointConfig(num_to_keep=2),
),
)
Ray Train 会自动：
在所有可用节点上启动工作节点，并在需要时在自动缩放的 Ray 集群中添加新节点。
在所有工作节点间进行数据分片。
无需对训练函数进行任何更改。同一个 train_func_per_worker 在 1 个 GPU 或 256 个 GPU 上运行效果相同。
本教程使用 DistributedDataParallel (DDP)，它会在每个 GPU 上复制整个模型。对于无法放入单个 GPU 的大型模型，你可以通过设置 prepare_model(parallel_strategy="fsdp") 切换到 FullyShardedDataParallel (FSDP)，以跨工作节点对参数、梯度和优化器状态进行分片。
异构集群：分离数据和训练资源#
由于 Ray Data 和 Ray Train 是独立的系统，它们不需要共享同一台机器。默认情况下，Ray Data 预处理和训练工作节点都在同一节点上运行。但是，你可以选择向集群添加仅 CPU 节点，Ray Data 会自动将预处理任务调度到这些节点上，从而使昂贵的 GPU 节点可以专注于训练。
当数据预处理成为瓶颈时，这非常有用。如果你发现 GPU 使用率较低，因为工作节点在等待数据，你可以向集群添加更便宜的仅 CPU 节点，Ray Data 会将预处理扩展到这些节点上。
有关更多信息，请参阅 配置数据摄取。
容错机制#
长期运行的分布式训练作业容易受到硬件故障的影响。这些包括硬件故障、网络中断或抢占。如果没有容错机制，任何这些事件都可能迫使你从零开始重新训练，浪费时间和计算资源。
Ray Train 具有自动处理这些故障的功能。当工作进程崩溃时，Ray Train 会在原地重启它并恢复训练。如果整个节点宕机，Ray Train 会提供一个替代节点并从最近的检查点恢复，这样只会丢失少量的进度。这使得中断训练作业并在稍后恢复变得非常实用。
要启用自动故障恢复，请在 RunConfig 中配置 FailureConfig。max_failures 参数控制 Ray Train 在放弃前容忍多少次连续故障：
from ray.train import FailureConfig
run_config = RunConfig(
storage_path="s3://my-bucket/ray-checkpoints",
failure_config=FailureConfig(max_failures=3),
checkpoint_config=CheckpointConfig(num_to_keep=2),
)
有关更多详细信息，请参阅 Ray Train 容错指南。
监控训练作业#
监控对于运行分布式训练至关重要。Ray 仪表盘显示实时指标，包括：
每个 epoch 的训练损失和验证指标
每个工作节点的 GPU 利用率和内存使用情况
数据加载吞吐量
工作节点状态和错误日志
要查看仪表盘，请打开 Ray 初始化后打印在日志中的链接。通常该链接为 https://:8265。
仪表盘允许你：
监控所有工作节点的训练进度
检查来自各个工作节点的日志
识别数据加载或通信瓶颈
查看每个工作节点的 CPU、GPU 和内存资源使用情况
通过详细的错误消息和堆栈跟踪来调试故障
有关更多信息，请参阅 Ray Train 监控文档。
结论#
在本教程中，你：
使用 Hugging Face Transformers 和 PyTorch 预训练了一个 GPT-2（约 1.24 亿参数）语言模型。
使用分布式流式传输的 Ray Data 加载并预处理了 Wikitext-103 数据集。
使用 Ray Train 的
TorchTrainer在 4 个 GPU 上运行了分布式训练，且仅对标准 PyTorch 训练循环进行了极少的更改。学习了如何保存和加载用于模型恢复的分布式检查点。
学习了如何通过更改
ScalingConfig和RunConfig来扩展到多节点集群。学习了异构集群如何让你在 CPU 节点上运行数据预处理，并在 GPU 节点上进行训练，从而优化成本和性能。
了解了 Ray Train 用于生产环境训练作业的容错机制。
使用 Ray 仪表盘监控了训练过程。
延伸阅读#
脚本运行总耗时：（1 分 57.424 秒）