Resuming from checkpoint: data loader state (dataset shuffling, batch index)

[vadimkantorov]

Hi!

How is represented state of data iteration in Trinity?

In Verl StatefulDataLoader is used (and in general classic DataLoader components): https://github.com/volcengine/verl/blob/313366fd85e95ad43d567a808dd647089723a255/verl/trainer/ppo/ray_trainer.py#L990-L997

How is iteration index / data permutation is represented in Trinity?

Is this state saved to checkpoint? or some metadata near the checkpoint?

Is this state being loaded/recovered when resuming from a checkpoint?

In general, we're lacking understanding of data prep infrastructure of Trinity (and how it relates to Verl's - if any of the Verl's data loading components are reused in Trinity)

Thanks!

[pan-x-c]

In Trinity-RFT, the Buffer module replaces the traditional DataLoader, and we do not use veRL's DataLoader components.

Different buffers have different read and write strategies. Generally, Explorer uses FILE type buffer as its data source, while Trainer uses QUEUE type buffer.

FILE buffer is similar to veRL's StatefulDataLoader. It reads data strictly in the order in which it appears in the file. Additionally, when the Explorer executes save_checkpoint, the current state of the buffer is saved to the checkpoint_root_dir/project/name/monitor folder. When the experiment is restarted, the corresponding file will be checked first to see if there is previous buffer state. You can find the corresponding code in cache manager and explorer

[vadimkantorov]

When the experiment is restarted, the corresponding file will be checked first to see if there is previous buffer state

Only for the FILE buffer? Or others like QUEUE / SQL too?

[pan-x-c]

SQL updates the usage count each time it reads data, which is finer-grained than FILE, and each update to the database is persistent and automatically inherited on the next startup.

QUEUE, on the other hand, is a temporary in-memory queue that lacks state preservation capabilities but supports backing up written data to a SQL database (pointing the path field to a SQLite database file). In practice, we generally use FILE as explorer_input to persist training progress, and QUEUE as trainer_input to temporarily save Experience data.

[vadimkantorov]

How would you recommend to implement / control data shuffling in your setup?

[pan-x-c]

Currently, we manually shuffle the data before training, and the order is determined after training starts.

[hiyuchang]

Now we have supported shuffling data by #326 .
An example is:

buffer:
  explorer_input:
    tasksets:
      - name: math
        storage_type: file
        path: /cpfs/data/tir/data/filter_all_correct/math
        task_selector:
          selector_type: random