File Handles released frequently during fully random access workloads

[bernardhan33]

Hi Team and @philipp-fischer ,

I'm testing my workload that enables fully random access with shuffle_buffer_size==1 and max_samples_per_sequence==1 and I noticed that the storage throughput was pretty bad due to repeated OpenFile which makes the file handles released and recreated for every micro-batch.

Digging into the code I bit I noticed a few places which seem to provide customization for this behavior. Namely,

1. itar_cache_size (code ref) which controls the size of the LRU cache that stores the file descriptors.
2. parallel_shard_iters (code ref) which constructs the itar_cache_size parameter.

Ideally in the fully random access we'd want all file handles to be kept open until the end of the training.

I'm using the basic Megatron-Energon code:

train_dataset = get_train_dataset(
      cfg.model.energon.path,
      batch_size=cfg.model.micro_batch_size,
      task_encoder=task_encoder,
      worker_config=worker_config,
      max_samples_per_sequence=max_samples_per_sequence,
      packing_buffer_size=None,
      shuffle_buffer_size=cfg.model.energon.shuffle_buffer_size,
      split_part='train',
  )
  data_loader = get_savable_loader(train_dataset, worker_config=worker_config)

Therefore, while I'm figuring out ways to pass these parameters through, I'd like to reach out to see if I'm on the right track and if you can help me configure these parameters quickly. Thanks!

[bernardhan33]

Ah I realized that even if they are indeed the configurations to change, these lines of code haven't made to the released package yet? At least I'm not seeing it in 5.2.0?

Anyways, we would really love to optimize these workloads starting by keeping the file handles open for all the tar files. Any ideas / suggestions? Thanks!

[philipp-fischer]

Hi!
One of the main reasons for why we slice up the shards and distribute them to workers in the way we do, is to trade off file I/O against randomness.

Typically, you neither want (1.) sequential reading of the data like standard WebDataset, because it's not random enough, and you usually don't want (2.) fully random access because it's slow to perform random access for normal file systems and disks.
Full random access for every sample means you cannot make use of the internal caches of the disks and the operating system, and instead you will seek around a lot. Also it's not practical to keep thousands of open file handles.

The solution to this is the trade-off that energon provides.
You can get a pretty decent randomness, by utilizing a large number of workers across many ranks, each of them opening multiple iterators on the dataset and each of them maintaining a shuffle buffer.
The samples from all those ranks will become a global batch that is sufficiently mixed.
If you use say 32 nodes, each with 8 GPUs and 2 workers each, that will give you 512 different sequences that the samples will originate from. This can be further multiplied by using parallel_shard_iters.
Btw. you should always use datasets that have been pre-shuffled and are not in a certain order.

However, maybe your use case is special and we need to better understand why you are looking to use max_samples_per_sequence=1

[bernardhan33]

Hi @philipp-fischer , thanks for the detailed description! Yes you are right about many things:

1. We do not want native WDS that only provides shuffling the samples with buffering. With simple experiments we have found that native WDS would iteratively fetch samples from the beginning of the tar files therefore insufficient randomness.
2. We have used datasets that have been pre-shuffled into the WDS and finally the Megatron-Energon compatiable formats.

We appreciate the Megatron Energon data loader for its implementation of additional sharding and indexing into the tar files to provide a balance between randomness and sequential I/Os. Once I study more about the codebase, I'd like to file a different issue confirming the I/O patterns with you.

For this issue specifically, we'd want to use max_samples_per_sequence=1 for the reason of creating a fully-random workload since a lot of SSD-based solutions (such as Local SSD in Google Cloud Platform) can support the fully-random workloads nicely without any further optimizations (e.g., additional num_workers). Additionally, Megatron Energon supports max_samples_per_sequence=1 and indeed produces a fully random workload as I can confirm from the sample sequence.

I'm still studying the codebase so please correct me for any mistakes below. As I understand the I/O patterns so far from a couple of experiments,

• Each worker participates in a round-robin fashion to fetch a batch_size of samples.
• Each worker shuffles between the parallel_shard_iters open shards within the worker and fetch samples from the open shards with the shuffling buffering logic defined in shuffle_buffer_dataset.py.

Therefore, with a max_samples_per_sequence=1, we can see the sharding information like

rank=0, worker=0: shard_range=[pretrain-0.tar[0, 1), pretrain-0.tar[1, 2), pretrain-0.tar[2, 3), ...<33>, pretrain-26.tar[1378, 1379), pretrain-26.tar[1379, 1380), pretrain-26.tar[1380, 1381)] sum(count)=53859

rank=0, worker=1: shard_range=[pretrain-26.tar[1381, 1382), pretrain-26.tar[1382, 1383), pretrain-26.tar[1383, 1384), ...<33>, pretrain-43.tar[2759, 2760), pretrain-43.tar[2760, 2761), pretrain-43.tar[2761, 2762)] sum(count)=53859

As expected, there are actually only 27 tar files but 53859 shards allocated to a worker.

Therefore, I had thought that the ITarReader implementation to be hopeful because it implements a LRU cache with the key being the tar file ID and the value being the file descriptors. Since there are only a limited number of files, even though we will do a lot of seeks, we get to keep the file handles without closing and re-opening the file for every single shard. However the code is not available in 5.2.0 for now. Did I understand this correctly?

Alternatively, any suggestions in optimizing the fully-random workloads by Megatron Energon? Note that we are not blocked on using Energon as the recommended settings of max_samples_per_sequence=100 and shuffle_buffer_size=100 work great. We just really love the fully-random workloads to guarantee complete randomness.

Thank you!

[philipp-fischer]

You are right that the recent version of energon introduced a new way of indexing that includes a LRU-like cache for open file handles.
Increasing that cache size could indeed solve the problem you are seeing.
Have you tried setting parallel_shard_iters to 27?

So far we haven't optimized for the "full randomness use case" at all, but I agree it's something that can be useful. We might include a dedicated code path for that case where we wouldn't even need to generate all these single-sample shard ranges. Each worker would then be allowed to randomly jump around in the shard ranges that are assigned to it. @lvoegtle

For now, you can try to workaround with parallel_shard_iters and of course contributions are welcome. So feel free to open a PR. If you're planning to, it might make sense to have an initial joint discussion about the architecture.

[bernardhan33]

Apologies for the late response -- I've been working on other things recently but I did test with building the source from a commit that made into the release 5.2.0 or 6.0.1 and confirmed that the file handles are kept open, which was exactly what we wished to see. Thank you!

Before I close this issue out, I did want to ask a qq -- how do we configure parallel_shard_iters from get_train_dataset or get_savable_loader? I did some quick experiments but did not find that setting it as one of kwargs can be "passed" through to "parallel_shard_iters" in the "BaseWebdatasetFactory" construction. Maybe I had some typos somewhere but would like to confirm. Thanks!

[philipp-fischer]

It should get passed all the way through from get_train_dataset.
So just do

get_train_dataset(
    parallel_shard_iters=123,
    ...
)