Make dataloader use another random generator

[thomasw21]

This is in order to start synchronizing the dropout across TP.

[stas00]

From https://pytorch.org/docs/stable/data.html#randomness-in-multi-process-data-loading:

By default, each worker will have its PyTorch seed set to base_seed + worker_id, where base_seed is a long generated by main process using its RNG (thereby, consuming a RNG state mandatorily) or a specified generator. However, seeds for other libraries may be duplicated upon initializing workers, causing each worker to return identical random numbers. (See this section in FAQ.).

How do you interpret this writeup?

1. Will it still do + worker_id and ours is 2 - which would make the generators different
2. or is the doc vague and actually it'll use the same generator for all workers?

[thomasw21]

1. Yes, we don't actually care about the seed of each worker since data loading should not be random.

The issue was base_seed is a long generated by main process using its RNG (thereby, consuming a RNG state mandatorily) Which means that tp_rank=0 would consume a a RNG state, while other ranks wouldn't.

[stas00]

tp_rank=0 would consume a a RNG state, while other ranks wouldn't.

why won't the other ranks?

I think here by "main process" it refers to the process that spawns the dataloader workers. Won't that main process be each rank's process with relationship to its workers? i.e. there are multiple main processes in this context.

Since I haven't investigated this - have you validated that this fix actually changes something? I assume you do, but you haven't described this in the OP, hence the asking.

[thomasw21]

why won't the other ranks?

Megatron-DeepSpeed/megatron/training.py

Lines 1170 to 1174 in 87a9dba

	if train_dataloader is not None:
	train_data_iterator = iter(train_dataloader) if dl_type == 'single' \
	else iter(cyclic_iter(train_dataloader))
	else:
	train_data_iterator = None

This line only runs on tp_rank = 0 if you go up the file a bit. The main process here will refer to tp_rank=0

Since I haven't investigated this - have you validated that this fix actually changes something?

Yes I was able to reproduce the discrepancy in torch_rng_state accross tp ranks, and this fixed it.

I assume you do, but you haven't described this in the OP, hence the asking.

My bad, since we talked about this on slack I didn't copy paste the thread.

[stas00]

Thank you for adding the missing sync, Thomas