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

在大规模分布式训练场景中，使用ShareGPT这类对话数据集配合DeepSpeed进行多卡训练时，最常见的问题是数据分片未对齐——明明每个GPU都加载了数据，但执行后却出现样本重复、进程间数据错乱，甚至某些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_worker不宜过大，以免引发文件句柄冲突。

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，需格外注意数据分片与激活检查点的交互。

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