DeepSpeed分布式训练中ShareGPT数据集的加载方法

在大规模分布式训练场景下，使用ShareGPT等对话数据集配合DeepSpeed进行多卡训练，最常见问题就是数据分片未能对齐——即便每个GPU都加载了数据，训练时依然会出现样本重复、进程间数据错乱，甚至个别显卡因无数据而陷入饥饿状态。此类问题的根源往往在于数据集未适配DistributedSampler的索引逻辑，或未妥善处理变长文本带来的动态性。以下直接给出三种经过验证的封装方案，分别适用于不同规模与数据组织形式。

先做第一个判断：如果你的ShareGPT数据总规模仍在几十GB级别，且可以一次性加载至内存中，那么最稳妥的方式就是走HuggingFace Datasets的原生路径，配合PyTorch的DistributedSampler实现均匀切分。无需额外折腾。

一、使用HuggingFace Datasets + DistributedSampler封装

该方案的核心思路是将ShareGPT JSONL文件转换为Dataset对象，再通过torch.utils.data.DistributedSampler实现跨进程的互斥切分，同时保留原始顺序语义。整个流程可概括为六步：

1. 安装必要依赖：pip install datasets torch transformers，这步无需多言。

2. 加载原始数据并构建Dataset对象：dataset = load_dataset("json", data_files={"train": "sharegpt_clean.jsonl"}, split="train")。注意data_files路径需指向实际文件。

3. 定义预处理函数，将对话结构统一为纯文本格式：def format_sharegpt(example): return {"text": "".join([f"### {msg['from']}: {msg['value']}" for msg in example["conversations"]])}。可根据业务需求调整分隔符，例如使用\n换行。

4. 应用映射并执行分词：tokenized_ds = dataset.map(format_sharegpt).map(lambda x: tokenizer(x["text"], truncation=True, max_length=2048), batched=True)。truncation参数和max_length应根据模型上下文长度设置，切勿超过实际支持的最大长度。

5. 初始化DistributedSampler：sampler = DistributedSampler(tokenized_ds, shuffle=True, drop_last=True)。shuffle可根据训练需求开启，但建议打开drop_last，以避免最后一个batch样本数不一致引发的梯度同步问题。

6. 构建DataLoader：dataloader = DataLoader(tokenized_ds, batch_size=4, sampler=sampler, num_workers=4)。num_workers可根据CPU核心数适当调大，但不宜超过16。

该方案运行后，每个GPU获取的样本是确定互斥的，且顺序与原始数据一致，非常适合调试与可复现性要求较高的场景。

二、自定义IterableDataset配合DeepSpeed的data_parallel配置

当ShareGPT数据量巨大（例如几百GB甚至上TB），且按日期或ID拆分成多个JSONL文件时，上述一次性加载全部数据至内存的方案就不再现实。此时需采用流式读取，并配合rank感知的路径选择，实现无状态、可恢复的数据供给。

具体做法是继承torch.utils.data.IterableDataset，重写__iter__方法。关键步骤如下：

1. 定义类时传入rank和world_size，仅保留当前进程负责的文件：class ShareGPTIterableDataset(IterableDataset): def __init__(self, file_list, rank, world_size): self.file_list = [f for i, f in enumerate(file_list) if i % world_size == rank]。这里使用取模方式平分文件列表，确保每个进程只读取属于自己的那份。

2. 在__iter__中按行解析JSONL并yield单条样本：for file_path in self.file_list: with open(file_path) as f: for line in f: yield json.loads(line)。注意此步骤应保持轻量，避免过多预处理。

3. 实例化数据集时传入当前进程的rank和world_size：ds = ShareGPTIterableDataset(glob.glob("sharegpt_*.jsonl"), dist.get_rank(), dist.get_world_size())。glob用于匹配所有分片文件。

4. 由于IterableDataset无法使用DistributedSampler，需直接构建DataLoader并禁用自动采样：dataloader = DataLoader(ds, batch_size=2, num_workers=2)。注意num_workers不宜过大，以免造成文件句柄冲突。

5. 最后在DeepSpeed配置中显式关闭自动采样，否则DeepSpeed会尝试为DataLoader附加一个Sampler，从而与IterableDataset冲突："data_efficiency": {"enabled": false}。该配置位于ZeRO优化对应的json文件中。

此方案的优点是内存占用极低，且天然支持断点续训——只需在保存checkpoint时记录各进程当前读取的行号或文件偏移量即可。

三、基于DeepSpeed的DataLoader Hook注入分片逻辑

第三种方法更为底层，直接绕过PyTorch的原生采样器，在DeepSpeed初始化阶段注入rank专属的数据路径和偏移量。它适用于已经预分片完成、且需要严格管控每张显卡token吞吐量的场景，例如配合梯度累积进行精细调参。

操作步骤如下：

1. 先手动或通过脚本将ShareGPT原始文件按world_size切分为独立文件：split -l 50000 sharegpt_full.jsonl sharegpt_part_。其中50000为每个分片包含的行数，可按实际显存调整。

2. 在init_process_group完成后，根据当前rank找到对应的分片文件：part_file = f"sharegpt_part_{dist.get_rank():02d}"。注意补零以保证排序正确。

3. 使用HuggingFace的load_dataset加载单个分片：local_ds = load_dataset("json", data_files=part_file, split="train")。这样每个进程仅读取一个文件，彻底避免数据交叉。

4. 调用deepspeed.initialize时通过training_data参数传入此本地数据集：model_engine, optimizer, _, _ = deepspeed.initialize(model=model, training_data=local_ds, ...)。注意在DeepSpeed最新版本中，该参数可能名为train_dataset，具体请查阅官方文档。

5. 最后确保DeepSpeed配置中不启用会干扰数据流的选项，例如partition_activations和stage3_gather_16bit_weights_on_model_save。典型的安全配置为：{"zero_optimization": {"stage": 2}}。如需使用stage3，则需特别注意数据分片与激活检查点的交互。

此方法虽然灵活，但需要手动管理分片文件，适合数据预处理已高度标准化的团队。若数据源为动态生成且每轮训练分片可能不同，则前两种方案更为省心。