[Trainer,Dataset/Feature] Streamable Dataset Cache implementation with lazy loading support

[KohakuBlueleaf]

TL;DR:

1. This PR propose a streaming friendly dataset cache format
2. This PR propose a lazy loading mechanism of 1
3. This PR modify nodes_train implementation based on 1 and 2

Intro

This draft pull request propose a new dataset cache format which is better than previous full pickle based dataset cache. While previous cache system have stuff like sharding so in usage it should not blow up user's system ram, but it is still not a best practice.

Therefore, this PR propose a new cache format which utilize a meta/header + tensor data seperation format. The data (conditions and latents) will be encoded as "pickled non tensor data + data structure skeleton" + "safetensors of all tensor data".

As safetensors natively support streaming or key based addressing. This design can achieve less sys ram usage with way better IO overhead and proper streaming of dataset reading. Further more I implemented a lazy loading system upon this streamable dataset impl, basically the latents and the conditions in the dataset can be loaded with only the meta/header part with corresponding safetensors keys info. And only really be loaded/realize when needed.

Implementation detail

conditions and latents are all python native container type with nested structure (may be nested) and the size of them are mainly from the size of the tensor object in them (latent tensor, text embedding tensor...)

Instead of pickle the whole list of them, I try to replace all the tensor in those containers become the key for loading from safetensors file. And put those tensor data into safetensors directly.
Sharding still exist bcuz the safetensors will be large.

[KohakuBlueleaf]

cc @alexisrolland @rattus128
Need suggestions/comments

[github-actions]

✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.
_{Posted by the CLA Assistant Lite bot.}

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds a lazy, shard-based training dataset format using safetensors files paired with pickled skeleton structures. New helpers split/rejoin tensors and realize lazy structures on demand. ResolutionBucket, SaveTrainingDataset, and LoadTrainingDataset are rewritten to produce and consume this new format, and two new nodes (RealizeLazyLatents, RealizeLazyConditionings) materialize lazy data. The training pipeline (nodes_train.py) is updated to stream and realize lazy latents/conditioning per-batch instead of requiring fully eager tensors. A node load-order adjustment ensures nodes_dataset.py loads before nodes_train.py.

Changes

Area	Changes
Dataset format	New tensor split/rejoin, safetensors header reading, LazyLatent/LazyConditioning/LazyBatchSamples/LazyCondEntry wrappers, RealizeRequired exception, _ShardReader
Dataset nodes	Rewrote ResolutionBucket, SaveTrainingDataset, LoadTrainingDataset; added RealizeLazyLatents, RealizeLazyConditionings
Training pipeline	Lazy-aware TrainSampler, batch conditioning realization, bucket/multi-res lazy latent handling, dtype propagation, TrainLoraNode wiring
Node loading	Reordered extras_files to load nodes_dataset.py before nodes_train.py

Sequence Diagram(s)

sequenceDiagram
    participant SaveTrainingDataset
    participant TensorSplitter
    participant ShardFiles

    SaveTrainingDataset->>TensorSplitter: split tensors per sample
    TensorSplitter-->>SaveTrainingDataset: skeleton + shard_tensors map
    SaveTrainingDataset->>ShardFiles: write shard.safetensors + skeleton.pkl

Loading

sequenceDiagram
    participant TrainLoraNode
    participant TrainSampler
    participant LazyLatent
    participant LazyConditioning

    TrainLoraNode->>TrainSampler: init with lazy_conds
    TrainSampler->>LazyLatent: realize_rows(indices)
    TrainSampler->>LazyConditioning: realize_entries(indices)
    LazyLatent-->>TrainSampler: materialized latent batch
    LazyConditioning-->>TrainSampler: realized conditioning

Loading

Related Issues: None specified in provided context.

Related PRs: None specified in provided context.

Suggested labels: enhancement, training, dataset

Suggested reviewers: ComfyUI training/dataset maintainers

---

Poem:

A rabbit hops through shards of light,
Tensors split, then rejoined tight.
Lazy latents wait in shade,
Realized only when the batch is made.
Dataset first, then training runs—
Hop, hop, hurray, the pipeline hums! 🐇

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly describes the main change: a streamable dataset cache with lazy loading support.
Description check	✅ Passed	The description is directly related to the PR and explains the new cache format, lazy loading, and training integration.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}