PyTorch Lightning实现nanoGPT多GPU训练大内存映射数据集加载故障排查

云上老周

2026-7-16

问题

在8卡多GPU环境下，使用PyTorch Lightning实现nanoGPT时，加载16GB内存映射格式的OpenWebText数据集过程卡住无法运行，但单GPU训练或数据集文件较小时加载正常。需要解决大内存映射数据集的多GPU加载问题。

环境

训练脚本运行在SLURM管理的集群中。

环境配置

module load openmpi
module load cuda/12.1

export MASTER_ADDR=`hostname`
export MASTER_PORT=12802
export NCCL_PROTO=simple
export FI_EFA_FORK_SAFE=1
export FI_LOG_LEVEL=1
export FI_EFA_USE_DEVICE_RDMA=1
export NCCL_DEBUG=info

export PYTHONFAULTHANDLER=1
export CUDA_LAUNCH_BLOCKING=0
export OMPI_MCA_mtl_base_verbose=1
export FI_EFA_ENABLE_SHM_TRANSFER=0
export FI_PROVIDER=efa
export FI_EFA_TX_MIN_CREDITS=64
export NCCL_TREE_THRESHOLD=0

Python包版本

datasets==2.13.1
einops==0.6.1
jsonargparse==4.21.1
lightning==2.0.2
numpy==1.24.3
setuptools==67.7.2
torch==2.0.1
torchinfo==1.8.0
tqdm==4.65.0
transformers==4.29.2
pydantic==1.10.8
tensorboard==2.13.0

实现代码

数据预处理脚本

import os
from dataclasses import dataclass
from typing import Optional, Union

import numpy as np
from datasets import load_dataset
from jsonargparse import CLI
from lightning import seed_everything
from tqdm.auto import tqdm
from transformers import AutoTokenizer


@dataclass
class Config:
    dataset_name_or_path: str
    tokenizer_name_or_path: str
    output_dir: str
    streaming: bool = False
    val_split: Optional[Union[int, float]] = None
    seed: int = 0
    num_proc: int = 4
    batch_size: int = 1_000


def cli_main(config: Config) -> None:
    # Setup: Reproducibility & Output directories
    seed_everything(config.seed)
    os.makedirs(config.output_dir, exist_ok=True)

    # DataLoading & Preparation
    dataset = load_dataset(config.dataset_name_or_path, streaming=config.streaming)

    # Split training dataset into training & validation sets
    if config.val_split and "val" not in dataset.keys():
        split_dataset = dataset["train"].train_test_split(
            test_size=config.val_split, seed=config.seed, shuffle=True
        )
        dataset["train"] = split_dataset.pop("train")
        dataset["val"] = split_dataset.pop("test")

    # Run tokenization
    tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path)

    # For packed datasets the eos token must be provided
    if tokenizer.eos_token is None:
        tokenizer.eos_token = tokenizer.pad_token

    def add_eos_batched(ids):
        return [
            x + [tokenizer.eos_token_id] if x[-1] != tokenizer.eos_token_id else x
            for x in ids
        ]

    def tokenize_examples(examples):
        ids = tokenizer(examples["text"]).input_ids

        # Optionally append eos token
        ids = add_eos_batched(ids)
        length = [len(x) for x in ids]

        return {"ids": ids, "len": length}

    dataset = dataset.map(
        tokenize_examples,
        remove_columns=dataset["train"].column_names,
        desc="Running Tokenization",
        batched=True,
        num_proc=config.num_proc,
    )

    # Save tokenized dataset
    for split_name, split in dataset.items():
        filepath = os.path.join(config.output_dir, f"{split_name}.bin")
        np_split = np.memmap(
            filepath, dtype=np.uint16, mode="w+", shape=(sum(split["len"]),)
        )

        num_samples = len(split["len"])
        num_batches = int(np.ceil(num_samples / config.batch_size))

        desc = f"Saving {split_name} dataset. Num samples: {num_samples}. Num batches: {num_batches}"
        offset = 0
        for batch_idx in tqdm(range(num_batches), desc=desc):
            batch = split.shard(
                num_shards=num_batches, index=batch_idx, contiguous=True
            ).with_format("numpy")
            batch = np.concatenate(batch["ids"])

            np_split[offset : offset + len(batch)] = batch
            offset += len(batch)

        np_split.flush()


if __name__ == "__main__":
    config = CLI(Config, as_positional=False)
    cli_main(config)

数据加载代码

import os
from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
from lightning import LightningDataModule
from torch.utils.data import DataLoader, Dataset


class PackedDataset(Dataset):
    def __init__(
        self, filepath: str, block_size: int = 1024, filemode: str = "r"
    ) -> None:
        super().__init__()

        self.filepath = filepath
        self.block_size = block_size

        self.data = np.memmap(filepath, dtype=np.uint16, mode=filemode)

    def __len__(self) -> int:
        return len(self.data) - self.block_size

    def __getitem__(self, index: int) -> torch.Tensor:
        return torch.as_tensor(
            self.data[index : index + self.block_size].astype(np.int64)
        )


@dataclass
class PackedDataModuleConfig:
    data_dir: str
    block_size: int
    filemode: str = "r"
    train_batch_size: Optional[int] = None
    val_batch_size: Optional[int] = None
    test_batch_size: Optional[int] = None
    num_workers: int = 4
    pin_memory: bool = False
    persistent_workers: bool = False


class PackedDataModule(LightningDataModule):
    def __init__(self, config: PackedDataModuleConfig) -> None:
        super().__init__()

        self.config = config
        self.splits = os.listdir(self.config.data_dir)

    def setup(self, stage: Optional[str] = None) -> None:
        if self.config.train_batch_size and "train.bin" in self.splits:
            self.train_dataset = PackedDataset(
                os.path.join(self.config.data_dir, "train.bin"),
                self.config.block_size,
                self.config.filemode,
            )

        if self.config.val_batch_size and "val.bin" in self.splits:
            self.val_dataset = PackedDataset(
                os.path.join(self.config.data_dir, "val.bin"),
                self.config.block_size,
                self.config.filemode,
            )

        if self.config.test_batch_size and "test.bin" in self.splits:
            self.test_dataset = PackedDataset(
                os.path.join(self.config.data_dir, "test.bin"),
                self.config.block_size,
                self.config.filemode,
            )

    def train_dataloader(self) -> Optional[DataLoader]:
        if hasattr(self, "train_dataset"):
            return self._create_dataloader(
                self.train_dataset, self.config.train_batch_size
            )

    def val_dataloader(self) -> Optional[DataLoader]:
        if hasattr(self, "val_dataset"):
            return self._create_dataloader(self.val_dataset, self.config.val_batch_size)

    def test_dataloader(self) -> Optional[DataLoader]:
        if hasattr(self, "test_dataset"):
            return self._create_dataloader(
                self.test_dataset, self.config.test_batch_size
            )

    def _create_dataloader(self, dataset: Dataset, batch_size: int) -> DataLoader:
        return DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=False,  # Causes the process to hang for large files
            num_workers=self.config.num_workers,
            pin_memory=self.config.pin_memory,
            persistent_workers=self.config.persistent_workers,
        )

解决方案

1. 修复多进程内存映射初始化冲突

多GPU训练时，每个数据加载子进程会共享主进程的内存映射对象，导致文件描述符冲突。修改PackedDataset，延迟内存映射初始化到第一次访问时，确保每个子进程独立打开文件：

class PackedDataset(Dataset):
    def __init__(
        self, filepath: str, block_size: int = 1024, filemode: str = "r"
    ) -> None:
        super().__init__()
        self.filepath = filepath
        self.block_size = block_size
        self.filemode = filemode
        self._data = None

    @property
    def data(self):
        if self._data is None:
            self._data = np.memmap(self.filepath, dtype=np.uint16, mode=self.filemode)
        return self._data

    def __len__(self) -> int:
        return len(self.data) - self.block_size

    def __getitem__(self, index: int) -> torch.Tensor:
        return torch.as_tensor(
            self.data[index : index + self.block_size].astype(np.int64)
        )

2. 配置分布式采样器

多GPU环境下必须使用DistributedSampler确保数据均匀分配到各个GPU，避免重复加载或数据缺失：

from torch.utils.data.distributed import DistributedSampler

def _create_dataloader(self, dataset: Dataset, batch_size: int) -> DataLoader:
    sampler = None
    # 自动检测分布式环境
    if self.trainer is not None and self.trainer.strategy.distributed:
        sampler = DistributedSampler(dataset, shuffle=True)
    
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=sampler is None,  # 非分布式环境才启用原生shuffle
        sampler=sampler,
        num_workers=self.config.num_workers,
        pin_memory=self.config.pin_memory,
        persistent_workers=self.config.persistent_workers,
        drop_last=True  # 避免最后一个batch大小不一致导致训练报错
    )

3. 优化内存映射与存储设置

只读数据集使用mode="r"而非"r+"，减少文件系统锁开销
确保所有节点的数据集路径指向共享存储（如Lustre/NFS），避免本地副本不一致
禁用FI_EFA_ENABLE_SHM_TRANSFER=0可能导致进程间共享内存通信受阻，可尝试改为FI_EFA_ENABLE_SHM_TRANSFER=1

4. 调整数据加载进程数

8卡环境下，建议设置num_workers=8（每卡1个进程）或num_workers=16（每卡2个进程），平衡数据加载速度与内存开销。同时保持persistent_workers=True减少进程重建的时间消耗。

5. 优化NCCL通信参数

移除NCCL_TREE_THRESHOLD=0，让NCCL自动选择最优通信算法
添加NCCL_IB_DISABLE=1（如果仅使用EFA通信）
检查SLURM任务的网络资源分配，确保节点间带宽充足无限制

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

PyTorch Lightning实现nanoGPT多GPU训练大内存映射数据集加载故障排查

环境 #

环境配置 #

Python包版本 #

实现代码 #

数据预处理脚本 #

数据加载代码 #

1. 修复多进程内存映射初始化冲突 #

2. 配置分布式采样器 #

3. 优化内存映射与存储设置 #

4. 调整数据加载进程数 #

5. 优化NCCL通信参数 #

环境