structuredconfig for train.py and eval.py

llmfoundry/data/dataloader.py

@@ -3,10 +3,9 @@
 
 """Dataloader builder utilities."""
 
-from typing import Union
+from typing import Any, Dict
 
 from composer import DataSpec
-from omegaconf import DictConfig
 from transformers import PreTrainedTokenizerBase
 
 from llmfoundry import registry
@@ -18,9 +17,9 @@
 
 
 def build_dataloader(
-    cfg: DictConfig,
+    cfg: Dict[str, Any],
     tokenizer: PreTrainedTokenizerBase,
-    device_batch_size: Union[int, float],
+    device_batch_size: int,
 ) -> DataSpec:
     """Builds a dataloader from a config.
 
@@ -30,14 +29,15 @@ def build_dataloader(
         device_batch_size (int): The size of the batches (number of examples)
             that the dataloader will produce.
     """
-    kwargs = {
-        'cfg': cfg,
+    name = cfg.pop('name')
+    kwargs: Dict[str, Any] = {
+        **cfg,
         'tokenizer': tokenizer,
         'device_batch_size': device_batch_size,
     }
 
     return construct_from_registry(
-        name=cfg.name,
+        name=name,
         registry=registry.dataloaders,
         partial_function=False,
         pre_validation_function=None,

llmfoundry/data/packing.py

@@ -3,12 +3,11 @@
 
 import logging
 import tempfile
-from typing import Callable, Dict, Iterable, List, Literal, Optional, Tuple
+from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Tuple
 
 import numpy as np
 import torch
 from composer.utils import dist
-from omegaconf import DictConfig
 from transformers import PreTrainedTokenizerBase
 
 log = logging.getLogger(__name__)
@@ -318,7 +317,7 @@ def pad_tensor(tensor: torch.Tensor, pad_value: int):
 
 
 def auto_packing_ratio(
-    dataloader_cfg: DictConfig,
+    dataloader_cfg: Dict[str, Any],
     tokenizer: PreTrainedTokenizerBase,
     device_batch_size: int,
     num_packing_ratios: int = 20,
@@ -352,20 +351,21 @@ def auto_packing_ratio(
     # Set the seed so that auto packing is deterministic.
     reproducibility.seed_all(0)
 
-    max_seq_len = dataloader_cfg.dataset.max_seq_len
     # If max_seq_len is very small, skip profiling and select packing ratio of 1.
+    dataset_config = dataloader_cfg['dataset']
+    max_seq_len = dataset_config.get('max_seq_len')
     if max_seq_len <= 100:
         return 1
 
     min_ratio = 1
     max_ratio = max_seq_len / 100
     profiling_results = profile_packing(
-        dataloader_cfg,
-        tokenizer,
-        min_ratio,
-        max_ratio,
-        num_packing_ratios,
-        device_batch_size,
+        dataloader_cfg=dataloader_cfg,
+        tokenizer=tokenizer,
+        min_ratio=min_ratio,
+        max_ratio=max_ratio,
+        num_packing_ratios=num_packing_ratios,
+        device_batch_size=device_batch_size,
     )
 
     # Obtain the maximum packing_ratio/minimum padding that has no waste.
@@ -392,7 +392,7 @@ def auto_packing_ratio(
 
 
 def profile_packing(
-    dataloader_cfg: DictConfig,
+    dataloader_cfg: Dict[str, Any],
     tokenizer: PreTrainedTokenizerBase,
     min_ratio: float,
     max_ratio: float,
@@ -416,39 +416,40 @@ def profile_packing(
 
     from llmfoundry.data.dataloader import build_dataloader
 
-    max_seq_len = dataloader_cfg.dataset.get('max_seq_len')
-    max_leftovers_to_keep = dataloader_cfg.dataset.get(
-        'max_leftovers_to_keep',
-        None,
-    )
+    dataset_cfg = dataloader_cfg['dataset']
+    max_seq_len = dataset_cfg.get('max_seq_len')
+    max_leftovers_to_keep = dataset_cfg.get('max_leftovers_to_keep', None)
 
     # Turn off packing and sequence parallelism for the dataloader (we want raw, pre-packed, full-length examples)
     dataloader_cfg = copy.deepcopy(dataloader_cfg)
-    dataloader_cfg.dataset.packing_ratio = 1.0
-    dataloader_cfg.dataset.auto_packing_replication = dataloader_cfg.dataset.get(
-        'seq_parallel_replication',
-        1,
-    ) or 1
-    dataloader_cfg.dataset.seq_parallel_replication = 1
-    dataloader_cfg.drop_last = False
-    dataloader_cfg.num_workers = 0
-    dataloader_cfg.prefetch_factor = None
-    dataloader_cfg.persistent_workers = False
+    dataloader_cfg.update({
+        'drop_last': False,
+        'num_workers': 0,
+        'prefetch_factor': None,
+        'persistent_workers': False,
+    })
+    dataloader_cfg['dataset']['packing_ratio'] = 1.0
+    dataloader_cfg['dataset']['auto_packing_replication'
+                             ] = dataloader_cfg['dataset'].get(
+                                 'seq_parallel_replication',
+                                 1,
+                             ) or 1
+    dataloader_cfg['dataset']['seq_parallel_replication'] = 1
 
     # If streaming dataset, use a temporary local folder for profiling
     local_rank_zero = dist.get_global_rank() - dist.get_local_rank()
-    if dataloader_cfg.dataset.get('remote') is not None:
+    if dataloader_cfg['dataset'].get('remote') is not None:
         tmp_path_to_broadcast = tempfile.TemporaryDirectory().name
         gathered_paths = dist.all_gather_object(tmp_path_to_broadcast)
         tmp_path = gathered_paths[local_rank_zero]
-        dataloader_cfg.dataset.local = tmp_path
+        dataloader_cfg['dataset']['local'] = tmp_path
 
-    if dataloader_cfg.dataset.get('streams') is not None:
-        for stream_config in dataloader_cfg.dataset.streams.values():
+    if dataloader_cfg['dataset'].get('streams') is not None:
+        for stream_config in dataloader_cfg['dataset']['streams'].values():
             tmp_path_to_broadcast = tempfile.TemporaryDirectory().name
             gathered_paths = dist.all_gather_object(tmp_path_to_broadcast)
             tmp_path = gathered_paths[local_rank_zero]
-            stream_config.local = tmp_path
+            stream_config['local'] = tmp_path
 
     # Determine the packing_ratio values we'll try
     packing_ratios, raw_batch_sizes = [], []

llmfoundry/data/text_data.py

@@ -20,8 +20,6 @@
 import numpy as np
 import torch
 from composer.core.data_spec import DataSpec
-from omegaconf import DictConfig
-from omegaconf import OmegaConf as om
 from streaming import Stream, StreamingDataset
 from torch.utils.data import DataLoader
 from transformers import PreTrainedTokenizerBase
@@ -268,79 +266,96 @@ def get_sequence_id_from_batch(
         return torch.cat([left_zeros, cumulative_sep[:, :-1]], dim=1)
 
 
-def build_streams(dataset_cfg: DictConfig):
-    streams_dict = dataset_cfg.pop('streams', None)
+def build_streams(streams: Optional[Dict[str, Any]] = None,):
+    streams_dict = streams
     # build streams
-    streams = None
+    streams_ret = []
     if streams_dict is not None:
-        streams = []
-        for stream in streams_dict.values():
-            # stream is the streams kwargs
-            # fwd all kwargs with **stream allows streaming to check args
-            streams.append(Stream(**stream))
-    return streams
+        streams_ret = [Stream(**stream) for stream in streams_dict.values()]
+    return streams_ret
 
 
 def build_text_dataloader(
-    cfg: DictConfig,
     tokenizer: PreTrainedTokenizerBase,
-    device_batch_size: Union[int, float],
+    device_batch_size: int,
+    dataset: Dict[str, Any],
+    drop_last: bool,
+    num_workers: int,
+    pin_memory: bool = True,
+    prefetch_factor: int = 2,
+    persistent_workers: bool = True,
+    timeout: int = 0,
 ) -> DataSpec:
-    assert cfg.name == 'text', f'Tried to build text dataloader with cfg.name={cfg.name}'
+
+    dataset_cfg = dataset
 
     # get kwargs
-    cfg.dataset['replication'], dataset_batch_size = construct_from_registry(
+    dataset_cfg['replication'], dataset_batch_size = construct_from_registry(
         name='dataset_replication_validator',
         registry=registry.dataset_replication_validators,
         partial_function=False,
         kwargs={
-            'cfg': cfg,
+            'cfg': dataset_cfg,
             'tokenizer': tokenizer,
             'device_batch_size': device_batch_size,
         },
     )
 
-    streams = build_streams(cfg.dataset)
+    streams = build_streams(
+        streams=dataset_cfg.pop('streams')
+        if 'streams' in dataset_cfg else None,
+    )
 
     valid_streaming_text_dataset_parameters = inspect.signature(
         StreamingTextDataset,
     ).parameters
 
     dataset_config_subset_for_streaming_text_dataset = {
         k: v
-        for k, v in cfg.dataset.items()
+        for k, v in dataset_cfg.items()
         if k in valid_streaming_text_dataset_parameters
     }
 
     # build dataset potentially with streams
-    dataset = StreamingTextDataset(
+    text_dataset = StreamingTextDataset(
         tokenizer=tokenizer,
         streams=streams,
         batch_size=dataset_batch_size,
         **dataset_config_subset_for_streaming_text_dataset,
     )
 
+    dataloader_cfg = {
+        'name': 'text',
+        'dataset': dataset_cfg,
+        'drop_last': drop_last,
+        'num_workers': num_workers,
+        'pin_memory': pin_memory,
+        'prefetch_factor': prefetch_factor,
+        'persistent_workers': persistent_workers,
+        'timeout': timeout,
+    }
+
     collate_fn, dataloader_batch_size = construct_from_registry(
         name='text_collator',
         registry=registry.collators,
         partial_function=False,
         kwargs={
-            'cfg': cfg,
-            'tokenizer': dataset.tokenizer,
+            'cfg': dataloader_cfg,
+            'tokenizer': tokenizer,
             'dataset_batch_size': dataset_batch_size,
         },
     )
 
     dl = DataLoader(
-        dataset,
+        text_dataset,
         collate_fn=collate_fn,
         batch_size=dataloader_batch_size,
-        drop_last=cfg.drop_last,
-        num_workers=cfg.num_workers,
-        pin_memory=cfg.get('pin_memory', True),
-        prefetch_factor=cfg.get('prefetch_factor', 2),
-        persistent_workers=cfg.get('persistent_workers', True),
-        timeout=cfg.get('timeout', 0),
+        drop_last=drop_last,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        prefetch_factor=prefetch_factor,
+        persistent_workers=persistent_workers,
+        timeout=timeout,
     )
 
     return construct_from_registry(
@@ -349,7 +364,7 @@ def build_text_dataloader(
         partial_function=False,
         kwargs={
             'dl': dl,
-            'dataset_cfg': cfg.dataset,
+            'dataset_cfg': dataset_cfg,
         },
     )
 
@@ -415,14 +430,17 @@ def build_text_dataloader(
         'drop_last': False,
         'num_workers': 4,
     }
-    cfg = om.create(cfg)
     device_batch_size = 2
 
     tokenizer_name = args.tokenizer
     tokenizer_kwargs = {'model_max_length': args.max_seq_len}
     tokenizer = build_tokenizer(tokenizer_name, tokenizer_kwargs)
 
-    loader = build_text_dataloader(cfg, tokenizer, device_batch_size).dataloader
+    loader = build_text_dataloader(
+        **cfg,
+        tokenizer=tokenizer,
+        device_batch_size=device_batch_size,
+    ).dataloader
     assert isinstance(loader, DataLoader)
     assert isinstance(loader.dataset, StreamingTextDataset)
     tokenizer = loader.dataset.tokenizer