WIP: add oscar slurm script for preprocess_data_dist

[adammoody]

I can't run on JZ, but for a concrete example, I think a script like the one in this PR could be used with the new preprocess_data_dist.py script. This requires the JSON support added in PR bigscience-workshop/Megatron-DeepSpeed#60

To process a JSON file, the script first generates an "index" that records the starting byte offset and length of each line in the source JSON file. That index file is stored beside the source JSON file to be reused in future runs. The index enables quick random access to the (variable-length) lines in the JSON file.

In the example SLURM script in this PR, for the source file:

$six_ALL_CCFRSCRATCH/datasets/oscar-small/oscar-en-shuffled-p1.jsonl

the preprocess_data_dist.py script will create the following files as a result of indexing the source JSON file:

$six_ALL_CCFRSCRATCH/datasets/oscar-small/oscar-en-shuffled-p1.jsonl.idx (persists after first run)
$six_ALL_CCFRSCRATCH/datasets/oscar-small/oscar-en-shuffled-p1.jsonl.idxtmp (created and deleted during run)

[adammoody]

If you verify this works well and correctly at smaller nodes counts, you might try scaling it higher. It has scaled well for me up to 64 nodes so far.

[adammoody]

Oh, and for a quick test, add something like preprocess_data_dist.py --count 1000 to limit the number of samples processed. It's good to test things with one or two nodes and a small sample count before trying large node counts and the full dataset.

[stas00]

Hi Adam,

I have just noticed your WIP PR here - is this still relevant and then let's merge it, or if not move/close it?

[adammoody]

Thanks, @stas00 . No need to merge this one.