如何基于Amazon SageMaker分布式创建Hugging Face Transformers分词器？

1. 数据分区与节点专属加载 #

先将大规模文本数据按节点数量做均匀分片（比如按文件哈希、前缀规则拆分），启动SageMaker Processor时，通过环境变量给每个节点分配唯一的NODE_INDEX和NUM_NODES标识。处理脚本中，每个节点仅加载对应自己分区的数据，避免全量数据重复加载：

import os
node_idx = int(os.environ.get("NODE_INDEX", 0))
num_nodes = int(os.environ.get("NUM_NODES", 1))
# 示例：加载S3上对应节点的分区数据
data_files = [f"s3://your-data-bucket/partition_{node_idx}/{f}" for f in os.listdir(f"s3://your-data-bucket/partition_{node_idx}/")]

2. 启用HuggingFace Tokenizers分布式训练 #

HuggingFace Tokenizers原生支持分布式训练，通过TrainingArguments开启分布式模式，让各节点同步词表更新，而非独立训练：

from tokenizers import Trainer, TrainingArguments
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer

# 初始化基础分词器模型
tokenizer = BPE()
trainer = BpeTrainer(special_tokens=["[CLS]", "[SEP]", "[PAD]"])

# 配置分布式训练参数
training_args = TrainingArguments(
    output_dir="./local_tokenizer",
    per_device_train_batch_size=128,
    distributed=True,  # 开启分布式同步
    gradient_accumulation_steps=2
)

# 绑定分区数据集并启动训练
trainer = Trainer(
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=your_partitioned_dataset
)
trainer.train()

3. SageMaker Processor分布式环境配置 #

创建HuggingFaceProcessor时，必须配置distribution参数启用Torch分布式通信，确保节点间能互通：

from sagemaker.huggingface import HuggingFaceProcessor

processor = HuggingFaceProcessor(
    role=your_iam_role,
    instance_count=4,  # 所需节点数
    instance_type="ml.c5.12xlarge",  # 或GPU实例如ml.p3.2xlarge
    transformers_version="4.30.0",
    tokenizers_version="0.13.3",
    distribution={
        "torch_distributed": {"enabled": True}
    }
)

# 启动处理任务，传递数据分区路径
processor.run(
    code="train_tokenizer.py",
    inputs=["s3://your-data-bucket/partitions/"],
    outputs=["s3://your-output-bucket/tokenizer/"]
)

4. 冗余流程规避 #

在训练脚本中添加节点角色判断：仅主节点（NODE_INDEX=0）负责最终的分词器保存、词表合并操作，其他节点完成分布式训练同步后直接退出，避免每个节点都重复生成完整分词器：

if node_idx == 0:
    # 主节点保存最终分词器
    tokenizer.save("./final_tokenizer.json")
    # 同步到S3
    s3.upload_file("./final_tokenizer.json", "your-output-bucket", "tokenizer/final_tokenizer.json")
else:
    exit(0)

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