feat: databricks deltalake dataset support #920
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Thanks for the PR! I tested this PR on a Databricks cluster and encountered a DeltaProtocolError. It appears the deltalake reader used here isn't compatible with tables that use Deletion Vectors or Column Mapping (which are now enabled by default on Databricks). Relevant discussions: delta-io/delta-rs#1094
|
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
akoumpa
force-pushed
the
akoumparouli/feat_streaming_deltalake
branch
from
3df865b to
0936f94
Compare
akoumpa
force-pushed
the
akoumparouli/feat_streaming_deltalake
branch
from
99fb5e1 to
b990139
Compare
akoumpa
force-pushed
the
akoumparouli/feat_streaming_deltalake
branch
from
1c75572 to
b8efd47
Compare
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
akoumpa
force-pushed
the
akoumparouli/feat_streaming_deltalake
branch
from
dd66497 to
caa8a6c
Compare
Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>
akoumpa
changed the title
Contributor
Author
|
/ok to test 937cf25 |
Contributor
Author
|
/ok to test 7ca3800 |
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
Contributor
Author
|
/ok to test 2ce5288 |
adil-a
reviewed
Jan 22, 2026
| global _DELTALAKE_AVAILABLE | ||
| if _DELTALAKE_AVAILABLE is None: | ||
| try: | ||
| import deltalake # noqa: F401 |
Collaborator
There was a problem hiding this comment.
can we clean up the noqa stuff
Contributor
Author
There was a problem hiding this comment.
the noqa is necessary otherwise ruff will complain
adil-a
reviewed
Jan 22, 2026
| ) | ||
|
|
||
|
|
||
| class HFDeltaLakeDataset: |
Collaborator
There was a problem hiding this comment.
this class doesn't implement the select method. same for DeltaLakeDataset. without this, the user cannot limit the number of data samples that are loaded
| self._shard_info = (num_shards, index) | ||
| return self | ||
|
|
||
| def shuffle(self, buffer_size: int = 1000, seed: Optional[int] = None) -> "HFDeltaLakeDataset": |
Collaborator
There was a problem hiding this comment.
i don't see shuffle being used in the iterator. how is shuffling happening?
jgerh
reviewed
Jan 23, 2026
Contributor
jgerh
left a comment
jgerh
left a comment
There was a problem hiding this comment.
Completed tech pubs review and provided some copyedits and suggestions.
| @@ -1,3 +1,3 @@ | |||
| # Dataset Overview: LLM and VLM Datasets in NeMo Automodel | |||
|
|
|||
| This page summarizes the datasets supported in NeMo Automodel for LLM and VLM and shows how to plug in your own datasets using Python functions or the YAML `_target_` mechanism. | |||
Contributor
There was a problem hiding this comment.
Suggested change
| This page summarizes the datasets supported in NeMo Automodel for LLMs and VLMs and shows how to plug in your own datasets using Python functions or the YAML `_target_` mechanism. |
|
|
||
| - **HellaSwag (completion SFT)** | ||
| - Wrapper: `nemo_automodel.components.datasets.llm.hellaswag.HellaSwag` | ||
| - Use case: single-turn completion style SFT where a prompt (ctx) is followed by a gold continuation (ending) |
Contributor
There was a problem hiding this comment.
Suggested change
| - Use case: single-turn, completion-style SFT where a prompt (context) is followed by a gold continuation |
| - Use case: multi-turn conversations and tool calling in OpenAI chat format | ||
| - Sources: local JSON/JSONL or Hugging Face Hub dataset ID | ||
| - Key args: | ||
| - `path_or_dataset_id`: path to local file(s) or HuggingFace dataset ID |
Contributor
There was a problem hiding this comment.
Suggested change
| - `path_or_dataset_id`: path to local file(s) or Hugging Face dataset ID |
| - Sources: local JSON/JSONL or Hugging Face Hub dataset ID | ||
| - Key args: | ||
| - `path_or_dataset_id`: path to local file(s) or HuggingFace dataset ID | ||
| - `tokenizer`: tokenizer instance (required. Must have chat template support) |
Contributor
There was a problem hiding this comment.
Suggested change
| - `tokenizer`: tokenizer instance (required; must have chat template support) |
| packed_sequence_size: 8192 # > 0 enables packing | ||
| split_across_pack: false | ||
| ``` | ||
| Use a collater that pads to an FP8-friendly multiple when training with FP8: |
Contributor
There was a problem hiding this comment.
Suggested change
| Use a collate function that pads to an FP8-friendly multiple when training with FP8: |
Comment on lines
7
to
9
| * Quick prototyping across diverse instruction datasets | ||
| * Schema flexibility without needing codebase changes | ||
| * Consistent field names for training loops, regardless of dataset source |
Contributor
There was a problem hiding this comment.
Suggested change
| * Consistent field names for training loops, regardless of dataset source | |
| - Quick prototyping across diverse instruction datasets | |
| - Schema flexibility without requiring code changes | |
| - Consistent field names for training loops, regardless of dataset source |
| break | ||
| ``` | ||
|
|
||
| :::note |
Contributor
There was a problem hiding this comment.
Suggested change
| :::note | |
| :::{note} |
|
|
||
| ### Multi-Node Slurm Configuration | ||
|
|
||
| :::note |
| ### Multi-Node Slurm Configuration | ||
|
|
||
| :::note | ||
| **Note for Multi-Node Training**: When using Hugging Face datasets in multi-node setups, you need shared storage accessible by all nodes. Set the `HF_DATASETS_CACHE` environment variable to point to a shared directory (e.g., `HF_DATASETS_CACHE=/shared/hf_cache`) in the yaml file as shown, to ensure all nodes can access the cached datasets. |
Contributor
There was a problem hiding this comment.
Suggested change
| **Multi-Node Training**: When using Hugging Face datasets in multi-node setups, you need shared storage accessible by all nodes. Set the `HF_DATASETS_CACHE` environment variable to point to a shared directory (e.g., `HF_DATASETS_CACHE=/shared/hf_cache`) in the yaml file as shown, to ensure all nodes can access the cached datasets. |
| answer_only_loss_mask=False, # compute loss over full sequence | ||
| ) | ||
|
|
||
| print(remote_ds[0].keys()) # {'context', 'question', 'answer'} |
Contributor
There was a problem hiding this comment.
Agent flagged this line: remote_ds is undefined in this local example, so the print would fail.
Suggested change
| print(local_ds[0].keys()) # {'question', 'answer'} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No description provided.