Fix for Squad Dataset Download

[rhmukundan]

Fix: Prevent Race Condition When Downloading SQuAD Dataset

Current Issue

When running with multiple GPUs per node and/or multiple nodes, downloading the SQuAD dataset can fail due to a race condition - multiple processes attempt to download the dataset simultaneously. This makes the dataset download to fail.

Fix

This fix ensures that the SQuAD dataset is downloaded in a separate SLURM job using only 1 node and 1 GPU per node. This prevents concurrent downloads and eliminates the race condition.

with run.Experiment(exp_name) as exp:
        if not SKIP_IMPORT:
            assert args.hf_token is not None, "HF token is required for importing checkpoint from HuggingFace"
            exp.add(*import_ckpt_experiment(executor, model(), source=f"hf://{HF_MODEL_URI}"))
            # Prepare SQuAD dataset before finetuning
            exp.add(*prepare_squad_dataset_experiment(executor, HF_MODEL_URI, seq_length=4096))

[guyueh1]

I have a few questions

• where are these added functions used?
• Is the purpose to avoid downloading squad from the network and use cached files?
• if yes, I think existing squad dataset should be able to use cached files if you properly set HF_HOME or HF_HUB envvar, but I haven't tried it recently, i can confirm later.

[rhmukundan]

@guyueh1 Thanks for your comment.
If you noticed, it is using the same function in nemo to download the dataset. Currently it is part of the overall job which means (num_nodes > 1 and num_gpus_per_node > 1) which causes the race condition which prevents the download and throws an error.
In order to fix this, I have created prepare_squad_dataset and prepare_squad_dataset_experiment functions which create a separate slurm job (1 Node and 1 GPU Per node) to avoid the race condition and download the dataset successfully. The functions used are pretty much what is part of the current nemo implementation (https://github.com/NVIDIA/NeMo/blob/577b7fb77f6767e23dec2926ba9b9d2a3722713b/scripts/performance/utils.py#L523) to get the squad dataset. What has changed is the way how it is downloaded (1 node and 1 GPU-per-node) via a separate slurm job similar to the checkpoint download.

This implementation will be used in finetune recipes like the following:

with run.Experiment(exp_name) as exp:
        if not SKIP_IMPORT:
            assert args.hf_token is not None, "HF token is required for importing checkpoint from HuggingFace"
            exp.add(*import_ckpt_experiment(executor, model(), source=f"hf://{HF_MODEL_URI}"))
            # Prepare SQuAD dataset before finetuning
            exp.add(*prepare_squad_dataset_experiment(executor, HF_MODEL_URI, seq_length=4096))

[guyueh1]

@rhmukundan I see. Can you provide an example of how to use it in run script?

[rhmukundan]

@guyueh1: I have pushed the Finetuning LLAMA4 Maverick recipe (that has the fix included) and also added the fix to LLAMA3 Finetuning 70b file.
cc: @ashbhandare

[guyueh1]

@malay-nagda can you review the added utility functions in util.py?

[malay-nagda]

LGTM

[github-actions]

[🤖]: Hi @rhmukundan 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully.

So it might be time to merge this PR or get some approvals.

//cc @chtruong814 @ko3n1g @pablo-garay @thomasdhc

[guyueh1]

@rhmukundan is this ready for merge?

[guyueh1]

LGTM

[guyueh1]

it's best to add a command-line argument, for instance standalone_dataset_preparation, for users to indicate they want this or the default way; that argument should default to false

[guyueh1]

This is not very clear; previously without this PR , are we not downloading the dataset from huggingface? My impression is it's still done somewhere in the dataset building process, just not explicitly, right? I think the difference is here you are separating it out as a new nemo-run experiment. If comment like this, users would think setting it to True without a local file will error out, but in reality it won't?

Could you further explain what different things are happening between False and True here;