only prepare data on rank 0

[conglongli]

Also change the filename hashing from hash() to hashlib.sha256() because of https://stackoverflow.com/questions/27522626/hash-function-in-python-3-3-returns-different-results-between-sessions

Tested with step 1 training_scripts/multi_node/run_66b.sh that PR works fine.

[puyuanOT]

I got OOM issue after checking out to this PR. It looks like this data preparing approach is consuming much large cpu memory than the previous implementation.

[conglongli]

I got OOM issue after checking out to this PR. It looks like this data preparing approach is consuming much large cpu memory than the previous implementation.

@puyuanOT It's very unlikely that this PR leads to more CPU memory consumption, because if you look at the code change, the only thing this PR does is reducing the data preparation to a single rank on each node. If running the same script with and without this PR really leads to different result (OOM with this PR), please share your error log and script so that we could try to reproduce and debug. Thanks.

[conglongli]

@puyuanOT Actually there might be one thing related to potential memory consumption increase. Could you try changing the "return torch.load(train_fname), torch.load(eval_fname)" at

DeepSpeedExamples/applications/DeepSpeed-Chat/training/utils/data/data_utils.py

Line 327 in ab4e2e5

return torch.load(train_fname), torch.load(eval_fname)

to "train_dataset, eval_dataset = torch.load(train_fname), torch.load(eval_fname)" and "return train_dataset, eval_dataset"? For users with only 1 GPU, this difference might lead to more memory consumption.

[puyuanOT]

@conglongli, thank you for your response! I'm fairly certain that this PR has resulted in significantly higher memory consumption, as I manually tested the script with and without the PR. To provide more details, without the PR, it consumed approximately 120GB for my four A10 GPUs; however, with the PR, the peak consumption exceeded 190GB.

I will give the fix you suggested a try.

[puyuanOT]

I figured this out.

In the previous implementation:

    fname = str(hash(fname))  # hash the file name to avoid too long file name
    train_fname = f"{output_path}/traindata_{fname}.pt"
    eval_fname = f"{output_path}/evaldata_{fname}.pt"

    cache_found = os.path.isfile(train_fname) and os.path.isfile(eval_fname)
    buf_create_cache = torch.ByteTensor([not cache_found]).cuda()
    torch.distributed.all_reduce(buf_create_cache)

    if buf_create_cache.item() == 0:
        return torch.load(train_fname), torch.load(eval_fname)

The code re-preprocesses the data every time the script is called due to the instability of str(hash(fname)), which returns a different hash each time it's called. Consequently, the code never executes the line torch.load(train_fname), torch.load(eval_fname). Interestingly, the data preprocessing consumes significantly less memory than loading the precomputed torch tensors, so preparing the data each time actually uses less memory than loading the pre-computed tensors.

Empirically, when setting --data_path Dahoas/rm-static Dahoas/full-hh-rlhf Dahoas/synthetic-instruct-gptj-pairwise yitingxie/rlhf-reward-datasets openai/webgpt_comparisons stanfordnlp/SHP and 'max_length=1024', the script consumes less than 200GB memory when using 4 GPUs (processes) with the old implementation, but it consumes more than 350 GB with the new implementation.

[puyuanOT]

To clarify, the memory consumption mentioned above are peak consumption.

[puyuanOT]

@conglongli Not sure if there is way to both avoid re-processing and high memory consumption.

[conglongli]

@puyuanOT Thanks for the investigation. We are currently working on switching to memmap-based data management which hopefully could greatly reduce CPU memory consumption #450. However, it'd be a complex transition thus will take some time. So you might want to first rollback to the old implementation.