feat: Implement lazy data loading for Dataset

[Ratish1]

this pr improves the Dataset class in miles/utils/data.py to support lazy data loading, filtering, and indexed access for large datasets. This addresses memory consumption issues, in SFT workloads, by avoiding the in-memory materialization of entire datasets. @fzyzcjy #226

• Modified read_file to use row-by-row iteration for JSONL and batched reading for Parquet.
• Refactored Dataset to build a lightweight index of valid sample locations on initialization, applying max_length filtering during this phase.
• Implemented __getitem__ to read and process only the requested sample on demand.

[fzyzcjy]

Too busy now, I will try to squeeze some time and have a review later :)

[Ratish1]

No problem

[Ratish1]

Hello @fzyzcjy, I wanted to follow up on this PR , but if you are busy its fine you can look at it later. Thanks.

[fzyzcjy]

After an internal discussion, we think the most ideal code may be somehow different. Therefore, may be need to wait a little bit.

[Ratish1]

Hello @fzyzcjy , I wanted to follow up on this. Thanks.

[lhoestq]

Hi ! Quentin from HF Datasets here :)

I was checking if I could help with anything, in particular if you are interested in loading larger than RAM datasets. E.g. the datasets library loads many formats like Parquet and JSON, makes tokenization easy, and works with larger than RAM datasets since the data is loaded with memory mapped Arrow files on disk. It also has streaming features to load datasets bigger than disk

[fzyzcjy]

@Ratish1 I am really too busy recently to check more details about this :/ @zhaochenyang20 do you want to find someone to review these?

@lhoestq looks great! Yes it would be great to handle huge datasets here, and at the same time keep code clean, b/c the framework supports SFT and can have big datasets.

[Ratish1]

Np @fzyzcjy , maybe I can work with @lhoestq and use the HF datasets library here

[zhaochenyang20]

@Ratish1 Have conflicts right now.

[PopSoda2002]

Hi, thanks for your effort and contribution! Nice work! Can I ask why do not choose memmap in python and can you share some comparison before and after?

[PopSoda2002]

Hi @Ratish1 , are you working with datasets from HF? It that, that will be nice

[Ratish1]

Hi @Ratish1 , are you working with datasets from HF? It that, that will be nice

No I'm not. I will refactor this code with the datasets library instead now. Thanks.

[PopSoda2002]

Hi, thanks for the effort, can you please show some results or comparison?

[Ratish1]

Hi, thanks for the effort, can you please show some results or comparison?

Hi, I haven't compared since I think datasets would be much faster for large datasets, but I ran a memory benchmark using a 500 MB dummy JSONL file (5 million rows) to verify the lazy loading of the current implementation. This is what I got

BENCHMARK RESULTS (Lazy Loading)
Initial Memory:   371.73 MiB
Peak Memory:      469.66 MiB
Memory Increment: 97.93 MiB

[PopSoda2002]

It looks good for me but I think it needs a more careful look!

[lhoestq]

I had a quick look and it looks good yo me, I just had one small nit:

[zhaochenyang20]

This is a comment. May not need to be done this time:

You are parsing JSON and building messages dynamically in every getitem call. While this saves RAM, it adds significant CPU overhead during the training loop.

If the JSON parsing is heavy, we might need to use hf_dataset.map() during init to pre-process these fields into a more efficient format (Arrow-native), rather than parsing raw strings on the fly.

[Ratish1]

Yes this is correct. But I have not done this change yet, since our primary goal for this PR was resolving the RAM spike during initialization, I believe this lazy approach is the safest first step. If we find that JSON parsing becomes a bottleneck for GPU throughput in future benchmarks, we can definitely change it to a .map() based pre-processing step to offload that work to the initialization phase.

[zhaochenyang20]

Reproduce/Verification commands:

docker pull radixark/miles:latest

docker run --rm --gpus all --ipc=host --shm-size=16g \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -it radixark/miles:latest /bin/bash

rm -rf /root/miles

git clone -b data-loading https://github.com/Ratish1/miles.git

cd miles

pip install -e .

hf download zai-org/GLM-Z1-9B-0414 --local-dir /data/ratish/GLM-Z1-9B-0414

hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /data/ratish/dapo-math-17k

source scripts/models/glm4-9B.sh

PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py ${MODEL_ARGS[@]} --hf-checkpoint /data/ratish/GLM-Z1-9B-0414 --save /data/ratish/GLM-Z1-9B-0414_torch_dist

export CUDA_DEVICE_MAX_CONNECTIONS=1
export CUDA_VISIBLE_DEVICES=5,7

PYTHONPATH=/root/Megatron-LM python train.py \
    ${MODEL_ARGS[@]} \
    --hf-checkpoint /data/ratish/GLM-Z1-9B-0414 \
    --load /data/ratish/GLM-Z1-9B-0414_torch_dist \
    --prompt-data /data/ratish/dapo-math-17k/dapo-math-17k.jsonl \
    --input-key prompt \
    --label-key label \
    --apply-chat-template \
    --use-miles-router \
    --num-proc 4 \
    --rollout-batch-size 16 \
    --global-batch-size 16 \
    --n-samples-per-prompt 1 \
    --num-rollout 1 \
    --colocate \
    --sglang-mem-fraction-static 0.8

[PopSoda2002]

LGTM! And I think adding this to a doc which will explain a little bit will be better!

[nanjiangwill]

@Ratish1 hi, i am wondering if lazy loading is still working during filtering ds = ds.filter(partial(_filter_func, **filter_kwargs), num_proc=num_proc, desc="Filtering invalid samples") (i.e. no peak memory consumption

[zhaochenyang20]

@Ratish1 hi, i am wondering if lazy loading is still working during filtering ds = ds.filter(partial(_filter_func, **filter_kwargs), num_proc=num_proc, desc="Filtering invalid samples") (i.e. no peak memory consumption

Hi, I've double-checked the underlying mechanism. ds.filter does indeed maintain the lazy loading behavior.

Under the hood, datasets processes the data in batches by reading them from the memory-mapped (mmap) storage into RAM. Once the filter function is executed, the Python objects are immediately released. It does not materialize the entire dataset in memory at once. Therefore, even during the filtering phase, memory consumption stays at a small constant level (determined by the batch_size and num_proc), avoiding any memory peaks that scale with the dataset size. This is exactly why integrating the datasets library is so beneficial for our use case.

[zhaochenyang20]

#372 (comment)

Sorry for my bad. This is reverted. We need more unit tests for dataset consistency. 😭