[RFC] Stateful dataloader and custom samplers for `SFTTrainer`

[SumanthRH]

Summary

The current SFTTrainer in SkyRL is pretty barebones: We load the dataset and store tokenized sequences as a list in memory, and sample directly from this list.

This RFC proposes a simple refactor for the SFT trainer to

1. Switch to torchdata.stateful_dataloader.StatefulDataLoader for train and eval data loading, matching the RL trainer's dataloader and checkpointing pattern.
2. Support customization for dataloading through stateful samplers that implement __iter__. In addition, state_dict, and load_state_dict can be implemented for easier resume.

Goals

• Replace manual list shuffling and slicing in SFTTrainer.train() with
  StatefulDataLoader, with ckpt and resume support
• Keep existing SFT tokenization and collation behavior unchanged.
• Support stateful custom samplers for ordered data, data mixing, and
  curriculum learning.

Non-Goals

• Changing the tokenization or data collation implementation

Proposed Trainer modifications

Add these fields to SFTTrainer:

self.train_dataloader = None
self.eval_dataloader = None

Add builder methods:

def build_train_dataloader(self, tokenized: list[dict]) -> StatefulDataLoader:
    ...

def build_eval_dataloader(self, eval_tokenized: list[dict]) -> StatefulDataLoader:
    ...

def build_train_sampler(self, tokenized: list[dict]) -> torch.utils.data.Sampler[int]:
    ...

The default train dataloader should be:

def collate_sft_examples(examples, collator, batch_size):
    return collator(examples, batch_size=batch_size)


collate_fn = functools.partial(
    collate_sft_examples,
    collator=self.collator,
    batch_size=self.sft_cfg.batch_size,
)

seeded_generator = torch.Generator()
seeded_generator.manual_seed(self.sft_cfg.seed)

self.train_dataloader = StatefulDataLoader(
    tokenized,
    batch_size=self.sft_cfg.batch_size,
    sampler=self.build_train_sampler(tokenized),
    collate_fn=collate_fn,
    drop_last=True,
    generator=seeded_generator,
    num_workers=self.sft_cfg.dataloader_num_workers,
    persistent_workers=self.sft_cfg.dataloader_persistent_workers,
    multiprocessing_context=(
        "spawn" if self.sft_cfg.dataloader_num_workers > 0 else None
    ),
)

SFT Config

Keep the existing SFTConfig.num_workers field for tokenization workers.
Add separate dataloader fields:

dataloader_num_workers: int = 0
dataloader_persistent_workers: bool = False
sampler: str = "random"  # "random", "sequential", or "custom"
sampler_class_path: Optional[str] = None #  import path to the custom sampler if any
sampler_kwargs: dict[str, Any] = field(default_factory=dict)

Example: Stateful Sequential Sampler

This is a simple ordered sampler.

class StatefulSequentialSampler(torch.utils.data.Sampler[int]):
    def __init__(self, data_source):
        self.data_source = data_source
        self.position = 0

    def __iter__(self):
        while self.position < len(self.data_source):
            idx = self.position
            self.position += 1
            yield idx
        self.position = 0

    def __len__(self):
        return len(self.data_source)

    def state_dict(self):
        ...

    def load_state_dict(self, state):
        ...

Note: state_dict and load_state_dict for samplers are optional for StatefulDataloader - the dataloader will fastforward by the number of samples yielded if no state is available.

Example: CurriculumLearningSampler

Curriculum learning can be implemented in multiple ways. Users could define a gradual fine-grained difficulty level and sample from low to high difficulty as training progress. The curriculum can also be coarse , where we group samples into difficulty levels , and train in order of difficulty level, but randomly sample within a difficulty level.

Example curriculum learning strategy where we have 3 difficulty levels of easy, medium and hard:

During training:

• first third of training: sample from easy subset;
• second third: sample from medium subset;
• final third: sample from hard subset.

Conceptual API:

class CurriculumLearningSampler:
    def __init__(
        self,
        lengths: List[int],
        num_samples=cfg.num_steps * cfg.batch_size,
        seed=cfg.seed
    ):
        ...

    def __iter__(self):
       # Pick the subset based on the current position / counter
       # iterate based on lengths of different subsets
        ...

    def state_dict(self):
        ...

    def load_state_dict(self, state):
        ...

Usage:

  easy_data = [s for s in data if s["difficulty"] == "easy"]
  medium_data = [s for s in data if s["difficulty"] == "medium"]
  hard_data = [s for s in data if s["difficulty"] == "hard"]
  train_dataset = ConcatDataset([easy_data, medium_data, hard_data])
  sampler = CurriculumSampler(
      lengths=[len(easy_data), len(medium_data), len(hard_data)],
      num_samples=cfg.num_steps * cfg.batch_size,
     seed=cfg.seed
  )

This will be added as an example sampler in examples/, and not as a native SkyRL sampler for now.

Data mixing example can be implemented in a similar way, maybe use the native WeightedRandomSampler.

Test Plan

Add CPU tests:

• default random dataloader order is deterministic for the same seed;
• random dataloader order differs for a different seed;
• sampler="sequential" yields examples in order;
• sequential sampler state_dict / load_state_dict resumes at the next sample;
• custom sampler state is included in the dataloader checkpoint state;
• data mixing and curriculum samplers resumes correctly.

GPU/e2e tests can remain unchanged initially.

Other considerations

Training loops using custom samplers should iterate once on global_step rather than looping epoch by epoch over the dataloader. This is primarily to avoid resetting dataloader state at epoch boundaries for custom samplers.