[Refactor] Rename datasets to prepare for multimodal datasets (#1916)

As titled, made the following change of names:

- torchtitan.datasets -> torchtitan.hf_datasets
- torchtitan.datasets.hf_datasets ->
torchtitan.hf_datasets.text_datasets
- build_hf_datasets -> build_text_datasets
- build_hf_validation_datasets -> build_text_validation_datasets

Author: wwwjn

tests/unit_tests/test_dataset_checkpointing.py

@@ -10,8 +10,8 @@
 from datasets import load_dataset
 from torchtitan.components.tokenizer import HuggingFaceTokenizer
 from torchtitan.config import ConfigManager
-from torchtitan.datasets import DatasetConfig
-from torchtitan.datasets.hf_datasets import build_hf_dataloader, DATASETS
+from torchtitan.hf_datasets import DatasetConfig
+from torchtitan.hf_datasets.text_datasets import build_text_dataloader, DATASETS
 
 
 class TestDatasetCheckpointing(unittest.TestCase):
@@ -72,7 +72,7 @@ def _build_dataloader(self, dataset_name, batch_size, seq_len, world_size, rank)
             ]
         )
 
-        return build_hf_dataloader(
+        return build_text_dataloader(
             tokenizer=tokenizer,
             dp_world_size=world_size,
             dp_rank=rank,

tests/unit_tests/test_train_spec.py

@@ -15,8 +15,8 @@
 from torchtitan.components.optimizer import build_optimizers, OptimizersContainer
 from torchtitan.components.tokenizer import build_hf_tokenizer
 from torchtitan.config import Optimizer as OptimizerConfig
-from torchtitan.datasets.hf_datasets import build_hf_dataloader
 from torchtitan.distributed.parallel_dims import ParallelDims
+from torchtitan.hf_datasets.text_datasets import build_text_dataloader
 from torchtitan.models.llama3 import parallelize_llama
 from torchtitan.protocols import BaseModelArgs, ModelProtocol
 from torchtitan.protocols.train_spec import (
@@ -82,7 +82,7 @@ def test_register_train_spec(self):
             pipelining_fn=None,
             build_optimizers_fn=build_optimizers,
             build_lr_schedulers_fn=build_lr_schedulers,
-            build_dataloader_fn=build_hf_dataloader,
+            build_dataloader_fn=build_text_dataloader,
             build_tokenizer_fn=build_hf_tokenizer,
             build_loss_fn=build_cross_entropy_loss,
         )
@@ -103,7 +103,7 @@ def test_optim_hook(self):
             pipelining_fn=None,
             build_optimizers_fn=fake_build_optimizers_with_hook,
             build_lr_schedulers_fn=build_lr_schedulers,
-            build_dataloader_fn=build_hf_dataloader,
+            build_dataloader_fn=build_text_dataloader,
             build_tokenizer_fn=build_hf_tokenizer,
             build_loss_fn=build_cross_entropy_loss,
         )

torchtitan/components/validate.py

@@ -14,8 +14,8 @@
 from torchtitan.components.metrics import MetricsProcessor
 from torchtitan.components.tokenizer import BaseTokenizer
 from torchtitan.config import JobConfig
-from torchtitan.datasets.hf_datasets import build_hf_validation_dataloader
 from torchtitan.distributed import ParallelDims, utils as dist_utils
+from torchtitan.hf_datasets.text_datasets import build_text_validation_dataloader
 from torchtitan.tools import utils
 from torchtitan.tools.logging import logger
 
@@ -62,7 +62,7 @@ def __init__(
         self.job_config = job_config
         self.parallel_dims = parallel_dims
         self.loss_fn = loss_fn
-        self.validation_dataloader = build_hf_validation_dataloader(
+        self.validation_dataloader = build_text_validation_dataloader(
             job_config=job_config,
             dp_world_size=dp_world_size,
             dp_rank=dp_rank,

torchtitan/experiments/flux/dataset/flux_dataset.py

@@ -22,11 +22,11 @@
 
 from torchtitan.components.tokenizer import BaseTokenizer
 from torchtitan.config import JobConfig
-from torchtitan.datasets import DatasetConfig
 from torchtitan.experiments.flux.dataset.tokenizer import (
     build_flux_tokenizer,
     FluxTokenizer,
 )
+from torchtitan.hf_datasets import DatasetConfig
 from torchtitan.tools.logging import logger
 
 

torchtitan/experiments/flux/tests/unit_tests/test_flux_dataloader.py

@@ -11,12 +11,12 @@
 from datasets import load_dataset
 
 from torchtitan.config import ConfigManager
-from torchtitan.datasets import DatasetConfig
 from torchtitan.experiments.flux.dataset.flux_dataset import (
     _cc12m_wds_data_processor,
     build_flux_dataloader,
     DATASETS,
 )
+from torchtitan.hf_datasets import DatasetConfig
 
 
 class TestFluxDataLoader(unittest.TestCase):

torchtitan/experiments/forge/example_train.py

@@ -18,8 +18,8 @@
 from torchtitan.components.tokenizer import build_hf_tokenizer
 from torchtitan.components.validate import build_validator
 from torchtitan.config import ConfigManager, JobConfig
-from torchtitan.datasets.hf_datasets import build_hf_dataloader
 from torchtitan.distributed import utils as dist_utils
+from torchtitan.hf_datasets.text_datasets import build_text_dataloader
 from torchtitan.tools import utils
 from torchtitan.tools.logging import init_logger, logger
 from torchtitan.tools.profiling import (
@@ -57,7 +57,7 @@ def __init__(self, job_config: JobConfig):
         self.tokenizer = build_hf_tokenizer(job_config)
 
         # build dataloader
-        self.dataloader = build_hf_dataloader(
+        self.dataloader = build_text_dataloader(
             dp_world_size=self.dp_degree,
             dp_rank=self.dp_rank,
             tokenizer=self.tokenizer,

torchtitan/experiments/simple_fsdp/deepseek_v3/__init__.py

@@ -8,8 +8,8 @@
 from torchtitan.components.lr_scheduler import build_lr_schedulers
 from torchtitan.components.optimizer import build_optimizers_with_moe_load_balancing
 from torchtitan.components.tokenizer import build_hf_tokenizer
-from torchtitan.datasets.hf_datasets import build_hf_dataloader
 from torchtitan.distributed.pipeline_parallel import pipeline_llm
+from torchtitan.hf_datasets.text_datasets import build_text_dataloader
 from torchtitan.models.deepseek_v3 import deepseekv3_args
 from torchtitan.protocols.train_spec import TrainSpec
 
@@ -25,7 +25,7 @@ def get_train_spec() -> TrainSpec:
         pipelining_fn=pipeline_llm,
         build_optimizers_fn=build_optimizers_with_moe_load_balancing,
         build_lr_schedulers_fn=build_lr_schedulers,
-        build_dataloader_fn=build_hf_dataloader,
+        build_dataloader_fn=build_text_dataloader,
         build_tokenizer_fn=build_hf_tokenizer,
         build_loss_fn=build_cross_entropy_loss,
     )

torchtitan/experiments/simple_fsdp/llama3/__init__.py

@@ -8,8 +8,8 @@
 from torchtitan.components.lr_scheduler import build_lr_schedulers
 from torchtitan.components.optimizer import build_optimizers
 from torchtitan.components.tokenizer import build_hf_tokenizer
-from torchtitan.datasets.hf_datasets import build_hf_dataloader
 from torchtitan.distributed.pipeline_parallel import pipeline_llm
+from torchtitan.hf_datasets.text_datasets import build_text_dataloader
 from torchtitan.models.llama3 import llama3_args
 from torchtitan.protocols.train_spec import TrainSpec
 
@@ -25,7 +25,7 @@ def get_train_spec() -> TrainSpec:
         pipelining_fn=pipeline_llm,
         build_optimizers_fn=build_optimizers,
         build_lr_schedulers_fn=build_lr_schedulers,
-        build_dataloader_fn=build_hf_dataloader,
+        build_dataloader_fn=build_text_dataloader,
         build_tokenizer_fn=build_hf_tokenizer,
         build_loss_fn=build_cross_entropy_loss,
     )

torchtitan/experiments/vlm/datasets/mm_datasets.py

@@ -22,7 +22,7 @@
 from torchtitan.components.dataloader import ParallelAwareDataloader
 from torchtitan.components.tokenizer import BaseTokenizer, HuggingFaceTokenizer
 from torchtitan.config import JobConfig
-from torchtitan.datasets import DatasetConfig
+from torchtitan.hf_datasets import DatasetConfig
 from torchtitan.tools.logging import logger
 
 from ..model.args import SpecialTokens
@@ -226,8 +226,8 @@ def _validate_mm_dataset(
     return path, config.loader, config.sample_processor
 
 
-class MultiModalDataset(IterableDataset, Stateful):
-    """MultiModal Dataset with support for sample packing."""
+class HuggingFaceMultiModalDataset(IterableDataset, Stateful):
+    """HuggingFace MultiModal Dataset with support for sample packing."""
 
     def __init__(
         self,
@@ -403,7 +403,7 @@ def build_mm_dataloader(
     packing_buffer_size = job_config.data.packing_buffer_size
     special_tokens = SpecialTokens.from_tokenizer(tokenizer)
 
-    dataset = MultiModalDataset(
+    dataset = HuggingFaceMultiModalDataset(
         dataset_name=job_config.training.dataset,
         dataset_path=dataset_path,
         tokenizer=tokenizer,